diff options
| author | Jeff Donahue <jeff.donahue@gmail.com> | 2014-08-21 12:53:43 -0700 |
|---|---|---|
| committer | Jeff Donahue <jeff.donahue@gmail.com> | 2014-08-21 13:08:25 -0700 |
| commit | 3b0f7a57f0ad6cd33ded7d148c7ad7e8f2a29e4d (patch) | |
| tree | bfa2844939fce7793834cf42fca0b27ee404f331 /tools/caffe.cpp | |
| parent | 43142a0e7a3c890b205d5fdc03f4e0c047c524ec (diff) | |
| download | caffeonacl-3b0f7a57f0ad6cd33ded7d148c7ad7e8f2a29e4d.tar.gz caffeonacl-3b0f7a57f0ad6cd33ded7d148c7ad7e8f2a29e4d.tar.bz2 caffeonacl-3b0f7a57f0ad6cd33ded7d148c7ad7e8f2a29e4d.zip | |
If specified, --gpu flag overrides SolverParameter solver_mode.
Diffstat (limited to 'tools/caffe.cpp')
| -rw-r--r-- | tools/caffe.cpp | 11 |
1 files changed, 11 insertions, 0 deletions
diff --git a/tools/caffe.cpp b/tools/caffe.cpp index bc404d3b..77e031ef 100644 --- a/tools/caffe.cpp +++ b/tools/caffe.cpp @@ -89,6 +89,17 @@ int train() { LOG(INFO) << "Starting Optimization"; caffe::SGDSolver<float> solver(solver_param); + + // Set device id and mode + if (FLAGS_gpu >= 0) { + LOG(INFO) << "Use GPU with device ID " << FLAGS_gpu; + Caffe::SetDevice(FLAGS_gpu); + Caffe::set_mode(Caffe::GPU); + } else if (!solver_param.has_solver_mode()) { + LOG(INFO) << "Use CPU."; + Caffe::set_mode(Caffe::CPU); + } + if (FLAGS_snapshot.size()) { LOG(INFO) << "Resuming from " << FLAGS_snapshot; solver.Solve(FLAGS_snapshot); |