1 #ifdef WITH_PYTHON_LAYER
2 #include "boost/python.hpp"
3 namespace bp = boost::python;
6 #include <gflags/gflags.h>
7 #include <glog/logging.h>
14 #include "boost/algorithm/string.hpp"
15 #include "caffe/caffe.hpp"
16 #include "caffe/util/signal_handler.h"
23 using caffe::shared_ptr;
27 using std::ostringstream;
29 DEFINE_string(gpu, "",
30 "Optional; run in GPU mode on given device IDs separated by ','."
31 "Use '-gpu all' to run on all available GPUs. The effective training "
32 "batch size is multiplied by the number of devices.");
33 DEFINE_string(solver, "",
34 "The solver definition protocol buffer text file.");
35 DEFINE_string(model, "",
36 "The model definition protocol buffer text file.");
37 DEFINE_string(phase, "",
38 "Optional; network phase (TRAIN or TEST). Only used for 'time'.");
39 DEFINE_int32(level, 0,
40 "Optional; network level.");
41 DEFINE_string(stage, "",
42 "Optional; network stages (not to be confused with phase), "
44 DEFINE_string(snapshot, "",
45 "Optional; the snapshot solver state to resume training.");
46 DEFINE_string(weights, "",
47 "Optional; the pretrained weights to initialize finetuning, "
48 "separated by ','. Cannot be set simultaneously with snapshot.");
49 DEFINE_int32(iterations, 50,
50 "The number of iterations to run.");
51 DEFINE_string(sigint_effect, "stop",
52 "Optional; action to take when a SIGINT signal is received: "
53 "snapshot, stop or none.");
54 DEFINE_string(sighup_effect, "snapshot",
55 "Optional; action to take when a SIGHUP signal is received: "
56 "snapshot, stop or none.");
58 // A simple registry for caffe commands.
59 typedef int (*BrewFunction)();
60 typedef std::map<caffe::string, BrewFunction> BrewMap;
63 #define RegisterBrewFunction(func) \
65 class __Registerer_##func { \
66 public: /* NOLINT */ \
67 __Registerer_##func() { \
68 g_brew_map[#func] = &func; \
71 __Registerer_##func g_registerer_##func; \
74 static BrewFunction GetBrewFunction(const caffe::string& name) {
75 if (g_brew_map.count(name)) {
76 return g_brew_map[name];
78 LOG(ERROR) << "Available caffe actions:";
79 for (BrewMap::iterator it = g_brew_map.begin();
80 it != g_brew_map.end(); ++it) {
81 LOG(ERROR) << "\t" << it->first;
83 LOG(FATAL) << "Unknown action: " << name;
84 return NULL; // not reachable, just to suppress old compiler warnings.
88 // Parse GPU ids or use all available devices
89 static void get_gpus(vector<int>* gpus) {
90 if (FLAGS_gpu == "all") {
93 CUDA_CHECK(cudaGetDeviceCount(&count));
97 for (int i = 0; i < count; ++i) {
100 } else if (FLAGS_gpu.size()) {
101 vector<string> strings;
102 boost::split(strings, FLAGS_gpu, boost::is_any_of(","));
103 for (int i = 0; i < strings.size(); ++i) {
104 gpus->push_back(boost::lexical_cast<int>(strings[i]));
107 CHECK_EQ(gpus->size(), 0);
111 // Parse phase from flags
112 caffe::Phase get_phase_from_flags(caffe::Phase default_value) {
113 if (FLAGS_phase == "")
114 return default_value;
115 if (FLAGS_phase == "TRAIN")
117 if (FLAGS_phase == "TEST")
119 LOG(FATAL) << "phase must be \"TRAIN\" or \"TEST\"";
120 return caffe::TRAIN; // Avoid warning
123 // Parse stages from flags
124 vector<string> get_stages_from_flags() {
125 vector<string> stages;
126 boost::split(stages, FLAGS_stage, boost::is_any_of(","));
130 // caffe commands to call by
131 // caffe <command> <args>
133 // To add a command, define a function "int command()" and register it with
134 // RegisterBrewFunction(action);
136 // Device Query: show diagnostic information for a GPU device.
138 LOG(INFO) << "Querying GPUs " << FLAGS_gpu;
141 for (int i = 0; i < gpus.size(); ++i) {
142 caffe::Caffe::SetDevice(gpus[i]);
143 caffe::Caffe::DeviceQuery();
147 RegisterBrewFunction(device_query);
149 // Translate the signal effect the user specified on the command-line to the
150 // corresponding enumeration.
151 caffe::SolverAction::Enum GetRequestedAction(
152 const std::string& flag_value) {
153 if (flag_value == "stop") {
154 return caffe::SolverAction::STOP;
156 if (flag_value == "snapshot") {
157 return caffe::SolverAction::SNAPSHOT;
159 if (flag_value == "none") {
160 return caffe::SolverAction::NONE;
162 LOG(FATAL) << "Invalid signal effect \""<< flag_value << "\" was specified";
165 // Train / Finetune a model.
167 CHECK_GT(FLAGS_solver.size(), 0) << "Need a solver definition to train.";
168 CHECK(!FLAGS_snapshot.size() || !FLAGS_weights.size())
169 << "Give a snapshot to resume training or weights to finetune "
171 vector<string> stages = get_stages_from_flags();
173 caffe::SolverParameter solver_param;
174 caffe::ReadSolverParamsFromTextFileOrDie(FLAGS_solver, &solver_param);
176 solver_param.mutable_train_state()->set_level(FLAGS_level);
177 for (int i = 0; i < stages.size(); i++) {
178 solver_param.mutable_train_state()->add_stage(stages[i]);
181 // If the gpus flag is not provided, allow the mode and device to be set
182 // in the solver prototxt.
183 if (FLAGS_gpu.size() == 0
184 && solver_param.has_solver_mode()
185 && solver_param.solver_mode() == caffe::SolverParameter_SolverMode_GPU) {
186 if (solver_param.has_device_id()) {
188 boost::lexical_cast<string>(solver_param.device_id());
189 } else { // Set default GPU if unspecified
190 FLAGS_gpu = "" + boost::lexical_cast<string>(0);
196 if (gpus.size() == 0) {
197 LOG(INFO) << "Use CPU.";
198 Caffe::set_mode(Caffe::CPU);
201 for (int i = 0; i < gpus.size(); ++i) {
202 s << (i ? ", " : "") << gpus[i];
204 LOG(INFO) << "Using GPUs " << s.str();
206 cudaDeviceProp device_prop;
207 for (int i = 0; i < gpus.size(); ++i) {
208 cudaGetDeviceProperties(&device_prop, gpus[i]);
209 LOG(INFO) << "GPU " << gpus[i] << ": " << device_prop.name;
212 solver_param.set_device_id(gpus[0]);
213 Caffe::SetDevice(gpus[0]);
214 Caffe::set_mode(Caffe::GPU);
215 Caffe::set_solver_count(gpus.size());
218 caffe::SignalHandler signal_handler(
219 GetRequestedAction(FLAGS_sigint_effect),
220 GetRequestedAction(FLAGS_sighup_effect));
222 if (FLAGS_snapshot.size()) {
223 solver_param.clear_weights();
224 } else if (FLAGS_weights.size()) {
225 solver_param.clear_weights();
226 solver_param.add_weights(FLAGS_weights);
229 shared_ptr<caffe::Solver<float> >
230 solver(caffe::SolverRegistry<float>::CreateSolver(solver_param));
232 solver->SetActionFunction(signal_handler.GetActionFunction());
234 if (FLAGS_snapshot.size()) {
235 LOG(INFO) << "Resuming from " << FLAGS_snapshot;
236 solver->Restore(FLAGS_snapshot.c_str());
239 LOG(INFO) << "Starting Optimization";
240 if (gpus.size() > 1) {
242 caffe::NCCL<float> nccl(solver);
243 nccl.Run(gpus, FLAGS_snapshot.size() > 0 ? FLAGS_snapshot.c_str() : NULL);
245 LOG(FATAL) << "Multi-GPU execution not available - rebuild with USE_NCCL";
250 LOG(INFO) << "Optimization Done.";
253 RegisterBrewFunction(train);
256 // Test: score a model.
258 CHECK_GT(FLAGS_model.size(), 0) << "Need a model definition to score.";
259 CHECK_GT(FLAGS_weights.size(), 0) << "Need model weights to score.";
260 vector<string> stages = get_stages_from_flags();
262 // Set device id and mode
265 if (gpus.size() != 0) {
266 LOG(INFO) << "Use GPU with device ID " << gpus[0];
268 cudaDeviceProp device_prop;
269 cudaGetDeviceProperties(&device_prop, gpus[0]);
270 LOG(INFO) << "GPU device name: " << device_prop.name;
272 Caffe::SetDevice(gpus[0]);
273 Caffe::set_mode(Caffe::GPU);
275 LOG(INFO) << "Use CPU.";
276 Caffe::set_mode(Caffe::CPU);
278 // Instantiate the caffe net.
279 Net<float> caffe_net(FLAGS_model, caffe::TEST, FLAGS_level, &stages);
280 caffe_net.CopyTrainedLayersFrom(FLAGS_weights);
281 LOG(INFO) << "Running for " << FLAGS_iterations << " iterations.";
283 vector<int> test_score_output_id;
284 vector<float> test_score;
286 for (int i = 0; i < FLAGS_iterations; ++i) {
288 const vector<Blob<float>*>& result =
289 caffe_net.Forward(&iter_loss);
292 for (int j = 0; j < result.size(); ++j) {
293 const float* result_vec = result[j]->cpu_data();
294 for (int k = 0; k < result[j]->count(); ++k, ++idx) {
295 const float score = result_vec[k];
297 test_score.push_back(score);
298 test_score_output_id.push_back(j);
300 test_score[idx] += score;
302 const std::string& output_name = caffe_net.blob_names()[
303 caffe_net.output_blob_indices()[j]];
304 LOG(INFO) << "Batch " << i << ", " << output_name << " = " << score;
308 loss /= FLAGS_iterations;
309 LOG(INFO) << "Loss: " << loss;
310 for (int i = 0; i < test_score.size(); ++i) {
311 const std::string& output_name = caffe_net.blob_names()[
312 caffe_net.output_blob_indices()[test_score_output_id[i]]];
313 const float loss_weight = caffe_net.blob_loss_weights()[
314 caffe_net.output_blob_indices()[test_score_output_id[i]]];
315 std::ostringstream loss_msg_stream;
316 const float mean_score = test_score[i] / FLAGS_iterations;
318 loss_msg_stream << " (* " << loss_weight
319 << " = " << loss_weight * mean_score << " loss)";
321 LOG(INFO) << output_name << " = " << mean_score << loss_msg_stream.str();
326 RegisterBrewFunction(test);
329 // Time: benchmark the execution time of a model.
331 CHECK_GT(FLAGS_model.size(), 0) << "Need a model definition to time.";
332 caffe::Phase phase = get_phase_from_flags(caffe::TRAIN);
333 vector<string> stages = get_stages_from_flags();
335 // Set device id and mode
338 if (gpus.size() != 0) {
339 LOG(INFO) << "Use GPU with device ID " << gpus[0];
340 Caffe::SetDevice(gpus[0]);
341 Caffe::set_mode(Caffe::GPU);
343 LOG(INFO) << "Use CPU.";
344 Caffe::set_mode(Caffe::CPU);
346 // Instantiate the caffe net.
347 Net<float> caffe_net(FLAGS_model, phase, FLAGS_level, &stages);
349 // Do a clean forward and backward pass, so that memory allocation are done
350 // and future iterations will be more stable.
351 LOG(INFO) << "Performing Forward";
352 // Note that for the speed benchmark, we will assume that the network does
353 // not take any input blobs.
355 caffe_net.Forward(&initial_loss);
356 LOG(INFO) << "Initial loss: " << initial_loss;
357 LOG(INFO) << "Performing Backward";
358 caffe_net.Backward();
360 const vector<shared_ptr<Layer<float> > >& layers = caffe_net.layers();
361 const vector<vector<Blob<float>*> >& bottom_vecs = caffe_net.bottom_vecs();
362 const vector<vector<Blob<float>*> >& top_vecs = caffe_net.top_vecs();
363 const vector<vector<bool> >& bottom_need_backward =
364 caffe_net.bottom_need_backward();
365 LOG(INFO) << "*** Benchmark begins ***";
366 LOG(INFO) << "Testing for " << FLAGS_iterations << " iterations.";
370 Timer backward_timer;
372 std::vector<double> forward_time_per_layer(layers.size(), 0.0);
373 std::vector<double> backward_time_per_layer(layers.size(), 0.0);
374 double forward_time = 0.0;
375 double backward_time = 0.0;
376 for (int j = 0; j < FLAGS_iterations; ++j) {
379 forward_timer.Start();
380 for (int i = 0; i < layers.size(); ++i) {
382 layers[i]->Forward(bottom_vecs[i], top_vecs[i]);
383 forward_time_per_layer[i] += timer.MicroSeconds();
385 forward_time += forward_timer.MicroSeconds();
386 backward_timer.Start();
387 for (int i = layers.size() - 1; i >= 0; --i) {
389 layers[i]->Backward(top_vecs[i], bottom_need_backward[i],
391 backward_time_per_layer[i] += timer.MicroSeconds();
393 backward_time += backward_timer.MicroSeconds();
394 LOG(INFO) << "Iteration: " << j + 1 << " forward-backward time: "
395 << iter_timer.MilliSeconds() << " ms.";
397 LOG(INFO) << "Average time per layer: ";
398 for (int i = 0; i < layers.size(); ++i) {
399 const caffe::string& layername = layers[i]->layer_param().name();
400 LOG(INFO) << std::setfill(' ') << std::setw(10) << layername <<
401 "\tforward: " << forward_time_per_layer[i] / 1000 /
402 FLAGS_iterations << " ms.";
403 LOG(INFO) << std::setfill(' ') << std::setw(10) << layername <<
404 "\tbackward: " << backward_time_per_layer[i] / 1000 /
405 FLAGS_iterations << " ms.";
408 LOG(INFO) << "Average Forward pass: " << forward_time / 1000 /
409 FLAGS_iterations << " ms.";
410 LOG(INFO) << "Average Backward pass: " << backward_time / 1000 /
411 FLAGS_iterations << " ms.";
412 LOG(INFO) << "Average Forward-Backward: " << total_timer.MilliSeconds() /
413 FLAGS_iterations << " ms.";
414 LOG(INFO) << "Total Time: " << total_timer.MilliSeconds() << " ms.";
415 LOG(INFO) << "*** Benchmark ends ***";
418 RegisterBrewFunction(time);
420 int main(int argc, char** argv) {
421 // Print output to stderr (while still logging).
422 FLAGS_alsologtostderr = 1;
424 gflags::SetVersionString(AS_STRING(CAFFE_VERSION));
426 gflags::SetUsageMessage("command line brew\n"
427 "usage: caffe <command> <args>\n\n"
429 " train train or finetune a model\n"
430 " test score a model\n"
431 " device_query show GPU diagnostic information\n"
432 " time benchmark model execution time");
433 // Run tool or show usage.
434 caffe::GlobalInit(&argc, &argv);
436 #ifdef WITH_PYTHON_LAYER
439 return GetBrewFunction(caffe::string(argv[1]))();
440 #ifdef WITH_PYTHON_LAYER
441 } catch (bp::error_already_set) {
447 gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/caffe");