SolverPool for MultiGPU in Python

include/caffe/parallel.hpp

@@ -3,6 +3,7 @@
 
 #include <boost/date_time/posix_time/posix_time.hpp>
 
+#include <string>
 #include <vector>
 
 #include "caffe/blob.hpp"
@@ -95,6 +96,12 @@ class P2PSync : public GPUParams<Dtype>, public Solver<Dtype>::Callback,
 
   void run(const vector<int>& gpus);
 
+  void set_up_gpus(const vector<int>& gpus);
+  void solve();
+  void step(int iters);
+  void tear_down();
+  vector<shared_ptr<Solver<Dtype> > > get_all_solvers();
+
  protected:
   void on_start();
   void on_gradients_ready();
@@ -106,13 +113,33 @@ class P2PSync : public GPUParams<Dtype>, public Solver<Dtype>::Callback,
   BlockingQueue<P2PSync<Dtype>*> queue_;
   const int initial_iter_;
   Dtype* parent_grads_;
-  shared_ptr<Solver<Dtype> > solver_;
+  boost::shared_ptr<Solver<Dtype> > solver_;
+  vector<shared_ptr<P2PSync<Dtype> > > syncs_;
 
   using Params<Dtype>::size_;
   using Params<Dtype>::data_;
   using Params<Dtype>::diff_;
 };
 
+// Synchronous data parallelism using map-reduce between local GPUs.
+template<typename Dtype>
+class SolverPool {
+ public:
+  SolverPool(string protofile, const vector<int>& gpus);
+  virtual ~SolverPool();
+
+  void Solve();
+
+  void Step(int iters);
+
+  vector<shared_ptr<Solver<Dtype> > > solvers();
+
+ protected:
+  shared_ptr<P2PSync<Dtype> > root_sync_;
+  shared_ptr<Solver<Dtype> > root_solver_;
+};
+
+
 }  // namespace caffe
 
 #endif

include/caffe/solver.hpp

@@ -39,6 +39,7 @@ class Solver {
     return test_nets_;
   }
   int iter() { return iter_; }
+  void set_device_id(int device) { param_.set_device_id(device); }
 
   // Invoked at specific points during an iteration
   class Callback {

python/caffe/_caffe.cpp

@@ -273,6 +273,14 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def("step", &Solver<Dtype>::Step)
     .def("restore", &Solver<Dtype>::Restore);
 
+  bp::class_<SolverPool<Dtype>,
+             shared_ptr<SolverPool<Dtype> >,
+             boost::noncopyable> (
+    "SolverPool", bp::init<string, vector<int> >())
+    .add_property("solvers", bp::make_function(&SolverPool<Dtype>::solvers,
+          bp::return_value_policy<bp::return_by_value>()))
+    .def("step", &SolverPool<Dtype>::Step);
+
   bp::class_<SGDSolver<Dtype>, bp::bases<Solver<Dtype> >,
     shared_ptr<SGDSolver<Dtype> >, boost::noncopyable>(
         "SGDSolver", bp::init<string>());
@@ -303,6 +311,9 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def(bp::vector_indexing_suite<vector<shared_ptr<Net<Dtype> > >, true>());
   bp::class_<vector<bool> >("BoolVec")
     .def(bp::vector_indexing_suite<vector<bool> >());
+  bp::class_<vector<shared_ptr<Solver<Dtype> > > >("SolverVec")
+    .def(bp::vector_indexing_suite<vector<shared_ptr<Solver<Dtype> > >
+                                   , true>());
 
   // boost python expects a void (missing) return value, while import_array
   // returns NULL for python3. import_array1() forces a void return value.

python/caffe/pycaffe.py

@@ -10,9 +10,10 @@
     from itertools import zip_longest as izip_longest
 import numpy as np
 
-from ._caffe import Net, SGDSolver
+from ._caffe import Net, SGDSolver, IntVec, SolverPool
 import caffe.io
 
+
 # We directly update methods from Net here (rather than using composition or
 # inheritance) so that nets created by caffe (e.g., by SGDSolver) will
 # automatically have the improved interface.
@@ -289,3 +290,10 @@ def _Net_batch(self, blobs):
 Net._batch = _Net_batch
 Net.inputs = _Net_inputs
 Net.outputs = _Net_outputs
+
+def get_solver_pool(solver_definition, gpus):
+  vec=caffe.IntVec()
+  vec.extend(gpus)
+  pool=caffe.SolverPool(solver_definition, gpus)
+  return pool
+

src/caffe/parallel.cpp

@@ -293,7 +293,7 @@ void P2PSync<Dtype>::on_start() {
 #ifdef DEBUG
   int device;
   CUDA_CHECK(cudaGetDevice(&device));
-  CHECK(device == solver_->param().device_id());
+  CHECK_EQ(device, solver_->param().device_id());
 #else
 //  CHECK(false);
 #endif
@@ -385,7 +385,15 @@ void P2PSync<Dtype>::on_gradients_ready() {
 
 template<typename Dtype>
 void P2PSync<Dtype>::run(const vector<int>& gpus) {
+  set_up_gpus(gpus);
+  solve();
+  tear_down();
+}
+
+template<typename Dtype>
+void P2PSync<Dtype>::set_up_gpus(const vector<int>& gpus) {
   // Pair devices for map-reduce synchronization
+  Caffe::SetDevice(gpus[0]);
   vector<DevicePair> pairs;
   DevicePair::compute(gpus, &pairs);
   ostringstream s;
@@ -394,16 +402,18 @@ void P2PSync<Dtype>::run(const vector<int>& gpus) {
   }
   LOG(INFO)<< "GPUs pairs " << s.str();
 
+  solver_->set_device_id(gpus[0]);
+
   SolverParameter param(solver_->param());
-  vector<shared_ptr<P2PSync<Dtype> > > syncs(gpus.size());
+  syncs_.resize(gpus.size());
 
   // Build the GPU tree by finding the parent for each solver
   for (int attempts = 0; attempts < pairs.size(); ++attempts) {
     for (int i = 1; i < pairs.size(); ++i) {
-      if (!syncs[i].get()) {
+      if (!syncs_[i].get()) {
         P2PSync<Dtype>* parent = NULL;
-        for (int j = 0; j < syncs.size(); ++j) {
-          P2PSync<Dtype>* sync = j == 0 ? this : syncs[j].get();
+        for (int j = 0; j < syncs_.size(); ++j) {
+          P2PSync<Dtype>* sync = j == 0 ? this : syncs_[j].get();
           if (sync) {
             const SolverParameter& p = sync->solver()->param();
             if (p.device_id() == pairs[i].parent()) {
@@ -413,29 +423,89 @@ void P2PSync<Dtype>::run(const vector<int>& gpus) {
         }
         if (parent) {
           param.set_device_id(pairs[i].device());
-          syncs[i].reset(new P2PSync<Dtype>(solver_, parent, param));
-          parent->children_.push_back((P2PSync<Dtype>*) syncs[i].get());
+          syncs_[i].reset(new P2PSync<Dtype>(solver_, parent, param));
+          parent->children_.push_back((P2PSync<Dtype>*) syncs_[i].get());
         }
       }
     }
   }
 
   LOG(INFO)<< "Starting Optimization";
 
-  for (int i = 1; i < syncs.size(); ++i) {
-    syncs[i]->StartInternalThread();
+  for (int i = 1; i < syncs_.size(); ++i) {
+    syncs_[i]->StartInternalThread();
   }
+}
+
+template<typename Dtype>
+void P2PSync<Dtype>::solve() {
+  LOG(INFO)<< "Starting Optimization";
 
   // Run root solver on current thread
   solver_->Solve();
+}
+
+template<typename Dtype>
+void P2PSync<Dtype>::tear_down() {
+  for (int i = 1; i < syncs_.size(); ++i) {
+    syncs_[i]->StopInternalThread();
+  }
+}
+
+template<typename Dtype>
+void P2PSync<Dtype>::step(int iters) {
+  // Run root solver on current thread
+  solver_->Step(iters);
+}
 
-  for (int i = 1; i < syncs.size(); ++i) {
-    syncs[i]->StopInternalThread();
+
+template<typename Dtype>
+SolverPool<Dtype>::SolverPool(string protofile,
+                        const vector<int>& gpus) {
+  Caffe::set_solver_count(gpus.size());
+  SolverParameter param;
+  ReadProtoFromTextFileOrDie(protofile, &param);
+  root_solver_.reset(GetSolver<Dtype>(param));
+
+  root_sync_.reset(
+      new caffe::P2PSync<Dtype>(root_solver_, NULL, root_solver_->param()));
+  root_sync_->set_up_gpus(gpus);
+}
+
+template<typename Dtype>
+void SolverPool<Dtype>::Step(int iters) {
+  root_sync_->step(iters);
+}
+
+template<typename Dtype>
+void SolverPool<Dtype>::Solve() {
+  root_sync_->solve();
+}
+
+template<typename Dtype>
+SolverPool<Dtype>::~SolverPool() {
+  root_sync_->tear_down();
+}
+
+template<typename Dtype>
+vector<shared_ptr<Solver<Dtype> > > SolverPool<Dtype>::solvers() {
+  return root_sync_->get_all_solvers();
+}
+
+template<typename Dtype>
+vector<shared_ptr<Solver<Dtype> > > P2PSync<Dtype>::get_all_solvers() {
+  vector<shared_ptr<Solver<Dtype> > > solver_list;
+  solver_list.push_back(solver_);
+  for (int i = 1; i < syncs_.size(); ++i) {
+    solver_list.push_back(shared_ptr<Solver<Dtype> >(syncs_[i]->solver()));
   }
+  return solver_list;
 }
 
 INSTANTIATE_CLASS(Params);
 INSTANTIATE_CLASS(GPUParams);
 INSTANTIATE_CLASS(P2PSync);
+INSTANTIATE_CLASS(SolverPool);
 
 }  // namespace caffe
+