1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
|
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <cstring>
#include <map>
#include <string>
#include <vector>
#include "caffe/caffe.hpp"
using caffe::Blob;
using caffe::Caffe;
using caffe::Net;
using caffe::Layer;
using caffe::shared_ptr;
using caffe::Timer;
using caffe::vector;
// Used in device query
DEFINE_int32(device_id, 0,
"[device_query,time] The device id to use.");
// Used in training
DEFINE_string(solver_proto_file, "",
"[train] The protobuf containing the solver definition.");
DEFINE_string(net_proto_file, "",
"[time] The net proto file to use.");
DEFINE_string(resume_point_file, "",
"[train] (optional) The snapshot from which to resume training.");
DEFINE_string(pretrained_net_file, "",
"[train] (optional) A pretrained network to run finetune from. "
"Cannot be set simultaneously with resume_point_file.");
DEFINE_int32(run_iterations, 50,
"[time] The number of iterations to run.");
DEFINE_bool(time_with_gpu, false,
"[time] Test the model with GPU.");
// A simple registry for caffe commands.
typedef int (*BrewFunction)();
typedef std::map<caffe::string, BrewFunction> BrewMap;
BrewMap g_brew_map;
#define RegisterBrewFunction(func) \
namespace { \
class __Registerer_##func { \
public: /* NOLINT */ \
__Registerer_##func() { \
g_brew_map[#func] = &func; \
} \
}; \
__Registerer_##func g_registerer_##func; \
}
static BrewFunction GetBrewFunction(const caffe::string& name) {
if (g_brew_map.count(name)) {
return g_brew_map[name];
} else {
LOG(ERROR) << "Available caffe actions:";
for (BrewMap::iterator it = g_brew_map.begin();
it != g_brew_map.end(); ++it) {
LOG(ERROR) << "\t" << it->first;
}
LOG(FATAL) << "Unknown action: " << name;
return NULL; // not reachable, just to suppress old compiler warnings.
}
}
// caffe actions that could be called in the form
// caffe.bin action
// To do so, define actions as "int action()" functions, and register it with
// RegisterBrewFunction(action);
int device_query() {
LOG(INFO) << "Querying device_id = " << FLAGS_device_id;
caffe::Caffe::SetDevice(FLAGS_device_id);
caffe::Caffe::DeviceQuery();
return 0;
}
RegisterBrewFunction(device_query);
int train() {
CHECK_GT(FLAGS_solver_proto_file.size(), 0);
caffe::SolverParameter solver_param;
caffe::ReadProtoFromTextFileOrDie(FLAGS_solver_proto_file, &solver_param);
LOG(INFO) << "Starting Optimization";
caffe::SGDSolver<float> solver(solver_param);
if (FLAGS_resume_point_file.size()) {
LOG(INFO) << "Resuming from " << FLAGS_resume_point_file;
solver.Solve(FLAGS_resume_point_file);
} else if (FLAGS_pretrained_net_file.size()) {
LOG(INFO) << "Finetuning from " << FLAGS_pretrained_net_file;
solver.net()->CopyTrainedLayersFrom(FLAGS_pretrained_net_file);
solver.Solve();
} else {
solver.Solve();
}
LOG(INFO) << "Optimization Done.";
return 0;
}
RegisterBrewFunction(train);
int time() {
// Set device id and mode
if (FLAGS_time_with_gpu) {
LOG(INFO) << "Use GPU with device id " << FLAGS_device_id;
Caffe::SetDevice(FLAGS_device_id);
Caffe::set_mode(Caffe::GPU);
} else {
LOG(INFO) << "Use CPU.";
Caffe::set_mode(Caffe::CPU);
}
// Instantiate the caffe net.
Caffe::set_phase(Caffe::TRAIN);
Net<float> caffe_net(FLAGS_net_proto_file);
// Do a clean forward and backward pass, so that memory allocation are done
// and future iterations will be more stable.
LOG(INFO) << "Performing Forward";
// Note that for the speed benchmark, we will assume that the network does
// not take any input blobs.
float initial_loss;
caffe_net.Forward(vector<Blob<float>*>(), &initial_loss);
LOG(INFO) << "Initial loss: " << initial_loss;
LOG(INFO) << "Performing Backward";
caffe_net.Backward();
const vector<shared_ptr<Layer<float> > >& layers = caffe_net.layers();
vector<vector<Blob<float>*> >& bottom_vecs = caffe_net.bottom_vecs();
vector<vector<Blob<float>*> >& top_vecs = caffe_net.top_vecs();
const vector<vector<bool> >& bottom_need_backward =
caffe_net.bottom_need_backward();
LOG(INFO) << "*** Benchmark begins ***";
LOG(INFO) << "Testing for " << FLAGS_run_iterations << " iterations.";
Timer total_timer;
total_timer.Start();
Timer forward_timer;
forward_timer.Start();
Timer timer;
for (int i = 0; i < layers.size(); ++i) {
const caffe::string& layername = layers[i]->layer_param().name();
timer.Start();
for (int j = 0; j < FLAGS_run_iterations; ++j) {
layers[i]->Forward(bottom_vecs[i], &top_vecs[i]);
}
LOG(INFO) << layername << "\tforward: " << timer.MilliSeconds() <<
" milli seconds.";
}
LOG(INFO) << "Forward pass: " << forward_timer.MilliSeconds() <<
" milli seconds.";
Timer backward_timer;
backward_timer.Start();
for (int i = layers.size() - 1; i >= 0; --i) {
const caffe::string& layername = layers[i]->layer_param().name();
timer.Start();
for (int j = 0; j < FLAGS_run_iterations; ++j) {
layers[i]->Backward(top_vecs[i], bottom_need_backward[i],
&bottom_vecs[i]);
}
LOG(INFO) << layername << "\tbackward: "
<< timer.MilliSeconds() << " milli seconds.";
}
LOG(INFO) << "Backward pass: " << backward_timer.MilliSeconds() <<
" milli seconds.";
LOG(INFO) << "Total Time: " << total_timer.MilliSeconds() <<
" milli seconds.";
LOG(INFO) << "*** Benchmark ends ***";
return 0;
}
RegisterBrewFunction(time);
int main(int argc, char** argv) {
caffe::GlobalInit(&argc, &argv);
CHECK_EQ(argc, 2);
return GetBrewFunction(caffe::string(argv[1]))();
}
|
