From e10fb59ccd3e3f1410936d1d1e5d415f013f9c00 Mon Sep 17 00:00:00 2001 From: Jeff Donahue Date: Fri, 9 May 2014 14:50:05 -0700 Subject: [PATCH] make solver_mode an enum with CPU and GPU -- fully backwards compatible with old 0/1 style --- docs/cifar10.md | 4 ++-- docs/mnist.md | 4 ++-- examples/cifar10/cifar10_full_solver.prototxt | 8 ++++---- examples/cifar10/cifar10_full_solver_lr1.prototxt | 4 ++-- examples/cifar10/cifar10_full_solver_lr2.prototxt | 4 ++-- examples/cifar10/cifar10_quick_solver.prototxt | 4 ++-- examples/cifar10/cifar10_quick_solver_lr1.prototxt | 4 ++-- examples/mnist/lenet_solver.prototxt | 4 ++-- examples/mnist/mnist_autoencoder_solver.prototxt | 3 ++- src/caffe/proto/caffe.proto | 6 +++++- src/caffe/solver.cpp | 3 ++- 11 files changed, 27 insertions(+), 21 deletions(-) diff --git a/docs/cifar10.md b/docs/cifar10.md index eb45e64..dd85667 100644 --- a/docs/cifar10.md +++ b/docs/cifar10.md @@ -89,7 +89,7 @@ CIFAR-10, while still small, has enough data to make GPU training attractive. To compare CPU vs. GPU training speed, simply change one line in all the `cifar*solver.prototxt`: - # solver mode: 0 for CPU and 1 for GPU - solver_mode: 0 + # solver mode: CPU or GPU + solver_mode: CPU and you will be using CPU for training. diff --git a/docs/mnist.md b/docs/mnist.md index 1e0f49f..9a9b46a 100644 --- a/docs/mnist.md +++ b/docs/mnist.md @@ -83,8 +83,8 @@ Um... How about GPU training? You just did! All the training was carried out on the GPU. In fact, if you would like to do training on CPU, you can simply change one line in `lenet_solver.prototxt`: - # solver mode: 0 for CPU and 1 for GPU - solver_mode: 0 + # solver mode: CPU or GPU + solver_mode: CPU and you will be using CPU for training. Isn't that easy? diff --git a/examples/cifar10/cifar10_full_solver.prototxt b/examples/cifar10/cifar10_full_solver.prototxt index cd62790..0a0b456 100644 --- a/examples/cifar10/cifar10_full_solver.prototxt +++ b/examples/cifar10/cifar10_full_solver.prototxt @@ -24,10 +24,10 @@ max_iter: 60000 # snapshot intermediate results snapshot: 10000 snapshot_prefix: "cifar10_full" -# solver mode: 0 for CPU and 1 for GPU +# solver mode: CPU or GPU # Note: there seems to be a bug with CPU computation in the pooling layers, -# and changing to solver_mode: 0 (for CPU computation) may result in NaNs -# on this example. If you want to train a variant of this architecture on the +# and changing to solver_mode: CPU may result in NaNs on this example. +# If you want to train a variant of this architecture on the # CPU, try changing the pooling regions from WITHIN_CHANNEL to ACROSS_CHANNELS # in both cifar_full_train.prototxt and cifar_full_test.prototxt. -solver_mode: 1 +solver_mode: GPU diff --git a/examples/cifar10/cifar10_full_solver_lr1.prototxt b/examples/cifar10/cifar10_full_solver_lr1.prototxt index 9f5f466..4376de5 100644 --- a/examples/cifar10/cifar10_full_solver_lr1.prototxt +++ b/examples/cifar10/cifar10_full_solver_lr1.prototxt @@ -24,5 +24,5 @@ max_iter: 65000 # snapshot intermediate results snapshot: 5000 snapshot_prefix: "cifar10_full" -# solver mode: 0 for CPU and 1 for GPU -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/examples/cifar10/cifar10_full_solver_lr2.prototxt b/examples/cifar10/cifar10_full_solver_lr2.prototxt index 785dffe..19580c5 100644 --- a/examples/cifar10/cifar10_full_solver_lr2.prototxt +++ b/examples/cifar10/cifar10_full_solver_lr2.prototxt @@ -24,5 +24,5 @@ max_iter: 70000 # snapshot intermediate results snapshot: 5000 snapshot_prefix: "cifar10_full" -# solver mode: 0 for CPU and 1 for GPU -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/examples/cifar10/cifar10_quick_solver.prototxt b/examples/cifar10/cifar10_quick_solver.prototxt index 32ba69d..4b547cc 100644 --- a/examples/cifar10/cifar10_quick_solver.prototxt +++ b/examples/cifar10/cifar10_quick_solver.prototxt @@ -23,5 +23,5 @@ max_iter: 4000 # snapshot intermediate results snapshot: 4000 snapshot_prefix: "cifar10_quick" -# solver mode: 0 for CPU and 1 for GPU -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/examples/cifar10/cifar10_quick_solver_lr1.prototxt b/examples/cifar10/cifar10_quick_solver_lr1.prototxt index 1f369cc..d4ba3d5 100644 --- a/examples/cifar10/cifar10_quick_solver_lr1.prototxt +++ b/examples/cifar10/cifar10_quick_solver_lr1.prototxt @@ -23,5 +23,5 @@ max_iter: 5000 # snapshot intermediate results snapshot: 5000 snapshot_prefix: "cifar10_quick" -# solver mode: 0 for CPU and 1 for GPU -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/examples/mnist/lenet_solver.prototxt b/examples/mnist/lenet_solver.prototxt index 50890d9..7947f2d 100644 --- a/examples/mnist/lenet_solver.prototxt +++ b/examples/mnist/lenet_solver.prototxt @@ -23,5 +23,5 @@ max_iter: 10000 # snapshot intermediate results snapshot: 5000 snapshot_prefix: "lenet" -# solver mode: 0 for CPU and 1 for GPU -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/examples/mnist/mnist_autoencoder_solver.prototxt b/examples/mnist/mnist_autoencoder_solver.prototxt index 9b48c30..06e057d 100644 --- a/examples/mnist/mnist_autoencoder_solver.prototxt +++ b/examples/mnist/mnist_autoencoder_solver.prototxt @@ -11,4 +11,5 @@ weight_decay: 0.0005 snapshot: 10000 snapshot_prefix: "mnist_autoencoder_train" momentum: 0.9 -solver_mode: 1 +# solver mode: CPU or GPU +solver_mode: GPU diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto index 9fba817..ab3c2fe 100644 --- a/src/caffe/proto/caffe.proto +++ b/src/caffe/proto/caffe.proto @@ -81,7 +81,11 @@ message SolverParameter { // debugging but the final protocol buffer size will be much larger. optional bool snapshot_diff = 16 [default = false]; // the mode solver will use: 0 for CPU and 1 for GPU. Use GPU in default. - optional int32 solver_mode = 17 [default = 1]; + enum SolverMode { + CPU = 0; + GPU = 1; + } + optional SolverMode solver_mode = 17 [default = GPU]; // the device_id will that be used in GPU mode. Use device_id = 0 in default. optional int32 device_id = 18 [default = 0]; // If non-negative, the seed with which the Solver will initialize the Caffe diff --git a/src/caffe/solver.cpp b/src/caffe/solver.cpp index 7442627..4932968 100644 --- a/src/caffe/solver.cpp +++ b/src/caffe/solver.cpp @@ -53,7 +53,8 @@ void Solver::Init(const SolverParameter& param) { template void Solver::Solve(const char* resume_file) { Caffe::set_mode(Caffe::Brew(param_.solver_mode())); - if (param_.solver_mode() && param_.has_device_id()) { + if (param_.solver_mode() == SolverParameter_SolverMode_GPU && + param_.has_device_id()) { Caffe::SetDevice(param_.device_id()); } Caffe::set_phase(Caffe::TRAIN); -- 2.7.4