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

mglaeser · cypof · commit 0d27efc7e3d3 · 2017-01-06T14:57:04.000-08:00
diff --git 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
@@ -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
@@ -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
@@ -44,10 +44,10 @@ def show_time():
         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:

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ class PythonLayer : public Layer<Dtype> {`
`21`	`21`	`// Disallow PythonLayer in MultiGPU training stage, due to GIL issues`
`22`	`22`	`// Details: https://github.com/BVLC/caffe/issues/2936`
`23`	`23`	`if (this->phase_ == TRAIN && Caffe::solver_count() > 1`
`24`		`- && !Caffe::root_solver() && !Caffe::multiprocess()) {`
	`24`	`+ && !Caffe::multiprocess()) {`
`25`	`25`	`LOG(FATAL) << "PythonLayer does not support CLI Multi-GPU, use train.py";`
`26`	`26`	`}`
`27`	`27`	`self_.attr("param_str") = bp::str(`