From: Jeff Donahue <jeff.donahue@gmail.com>
Date: Thu, 21 Aug 2014 19:53:43 +0000 (-0700)
Subject: If specified, --gpu flag overrides SolverParameter solver_mode.
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=ac85f8136b287f231ea495ecba5a72341018a1ae;p=platform%2Fupstream%2Fcaffe.git

If specified, --gpu flag overrides SolverParameter solver_mode.
---

diff --git a/tools/caffe.cpp b/tools/caffe.cpp
index 13398ae..eb32128 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);