diff options
| author | Evan Shelhamer <shelhamer@imaginarynumber.net> | 2017-01-17 09:12:01 -0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2017-01-17 09:12:01 -0800 |
| commit | 536851c00e44d545649648b372503f52a89c0499 (patch) | |
| tree | d8f841db4fec3a92bf2cbe9a651b2bdf2d29159f /tools/caffe.cpp | |
| parent | 0a91794d85562eaa6064541a58a3588b5afb0e5f (diff) | |
| parent | 5f28eb1147c1abb6e5e5c7cd282218679b0d531d (diff) | |
| download | caffeonacl-536851c00e44d545649648b372503f52a89c0499.tar.gz caffeonacl-536851c00e44d545649648b372503f52a89c0499.tar.bz2 caffeonacl-536851c00e44d545649648b372503f52a89c0499.zip | |
Merge pull request #4563 from cypof/nccl
adopt NVIDIA's NCCL for multi-GPU and switch interface to python
Diffstat (limited to 'tools/caffe.cpp')
| -rw-r--r-- | tools/caffe.cpp | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/tools/caffe.cpp b/tools/caffe.cpp index 9bf4214a..3587d8aa 100644 --- a/tools/caffe.cpp +++ b/tools/caffe.cpp @@ -195,6 +195,7 @@ int train() { // If the gpus flag is not provided, allow the mode and device to be set // in the solver prototxt. if (FLAGS_gpu.size() == 0 + && solver_param.has_solver_mode() && solver_param.solver_mode() == caffe::SolverParameter_SolverMode_GPU) { if (solver_param.has_device_id()) { FLAGS_gpu = "" + @@ -244,11 +245,15 @@ int train() { CopyLayers(solver.get(), FLAGS_weights); } + LOG(INFO) << "Starting Optimization"; if (gpus.size() > 1) { - caffe::P2PSync<float> sync(solver, NULL, solver->param()); - sync.Run(gpus); +#ifdef USE_NCCL + caffe::NCCL<float> nccl(solver); + nccl.Run(gpus, FLAGS_snapshot.size() > 0 ? FLAGS_snapshot.c_str() : NULL); +#else + LOG(FATAL) << "Multi-GPU execution not available - rebuild with USE_NCCL"; +#endif } else { - LOG(INFO) << "Starting Optimization"; solver->Solve(); } LOG(INFO) << "Optimization Done."; |