Python layers should build on multiprocess & solver_cnt; enable with bindings

diff --git a/include/caffe/layers/python_layer.hpp b/include/caffe/layers/python_layer.hpp
index 529b09c..10c4bfd 100644 (file)
--- a/include/caffe/layers/python_layer.hpp
+++ b/include/caffe/layers/python_layer.hpp
@@ -21,7 +21,7 @@ class PythonLayer : public Layer<Dtype> {
     // Disallow PythonLayer in MultiGPU training stage, due to GIL issues
     // Details: https://github.com/BVLC/caffe/issues/2936
     if (this->phase_ == TRAIN && Caffe::solver_count() > 1
-        && !Caffe::root_solver() && !Caffe::multiprocess()) {
+        && !Caffe::multiprocess()) {
       LOG(FATAL) << "PythonLayer does not support CLI Multi-GPU, use train.py";
     }
     self_.attr("param_str") = bp::str(

diff --git a/python/caffe/__init__.py b/python/caffe/__init__.py
index dde2e98..43a0c49 100644 (file)
--- a/python/caffe/__init__.py
+++ b/python/caffe/__init__.py
@@ -1,5 +1,5 @@
 from .pycaffe import Net, SGDSolver, NesterovSolver, AdaGradSolver, RMSPropSolver, AdaDeltaSolver, AdamSolver, NCCL, Timer
-from ._caffe import init_log, log, set_mode_cpu, set_mode_gpu, set_device, Layer, get_solver, layer_type_list, set_random_seed, solver_count, set_solver_count, solver_rank, set_solver_rank, Layer, get_solver
+from ._caffe import init_log, log, set_mode_cpu, set_mode_gpu, set_device, Layer, get_solver, layer_type_list, set_random_seed, solver_count, set_solver_count, solver_rank, set_solver_rank, set_multiprocess, Layer, get_solver
 from ._caffe import __version__
 from .proto.caffe_pb2 import TRAIN, TEST
 from .classifier import Classifier

diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
index 04dac23..3589e47 100644 (file)
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
@@ -359,6 +359,7 @@ BOOST_PYTHON_MODULE(_caffe) {
   bp::def("set_solver_count", &Caffe::set_solver_count);
   bp::def("solver_rank", &Caffe::solver_rank);
   bp::def("set_solver_rank", &Caffe::set_solver_rank);
+  bp::def("set_multiprocess", &Caffe::set_multiprocess);
 
   bp::def("layer_type_list", &LayerRegistry<Dtype>::LayerTypeList);
 

diff --git a/python/train.py b/python/train.py
index 730dbe7..5897f5d 100644 (file)
--- a/python/train.py
+++ b/python/train.py
@@ -44,10 +44,10 @@ def time(solver, nccl):
         if solver.iter % display == 0:
             s = '\n'
             for i in range(len(solver.net.layers)):
-                s += 'forw %3d %8s ' % (i, solver.net.layers[i].layer_param.name)
+                s += 'forw %3d %8s ' % (i, solver.net._layer_names[i])
                 s += ': %.2f\n' % fprop[i].ms
             for i in range(len(solver.net.layers) - 1, -1, -1):
-                s += 'back %3d %8s ' % (i, solver.net.layers[i].layer_param.name)
+                s += 'back %3d %8s ' % (i, solver.net._layer_names[i])
                 s += ': %.2f\n' % bprop[i].ms
             s += 'solver total: %.2f\n' % total.ms
             s += 'allreduce: %.2f\n' % allrd.ms
@@ -67,6 +67,7 @@ def solve(proto, snapshot, gpus, timing, uid, rank):
     caffe.set_device(gpus[rank])
     caffe.set_solver_count(len(gpus))
     caffe.set_solver_rank(rank)
+    caffe.set_multiprocess(True)
 
     solver = caffe.SGDSolver(proto)
     if snapshot and len(snapshot) != 0: