Merge pull request BVLC#961 from jeffdonahue/gpu-flag-overrides-solve…

…r-mode If specified, --gpu flag overrides SolverParameter solver_mode.
mitmul · Aug 21, 2014 · 5436ade · 5436ade
2 parents 12ba957 + 6413808
commit 5436ade
Showing 1 changed file with 11 additions and 0 deletions.
diff --git 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);