More detailed net_speed_benchmark

diff --git a/examples/net_speed_benchmark.cpp b/examples/net_speed_benchmark.cpp
index a287591..dd6f3ed 100644 (file)
--- a/examples/net_speed_benchmark.cpp
+++ b/examples/net_speed_benchmark.cpp
@@ -18,11 +18,35 @@
 using namespace caffe;
 
 int main(int argc, char** argv) {
-  cudaSetDevice(0);
-  Caffe::set_mode(Caffe::GPU);
-  Caffe::set_phase(Caffe::TRAIN);
-  int repeat = 100;
 
+  int total_iter = 50;
+
+  if (argc < 2) {
+    LOG(ERROR) << "net_speed_benchmark net_proto [iterations=50] [CPU/GPU] [Device_id=0]";
+    return 0;
+  }
+
+  if (argc >=3) {
+    total_iter = atoi(argv[2]);
+  }
+
+  LOG(ERROR) << "Testing for " << total_iter << "Iterations.";
+
+  if (argc >= 4 && strcmp(argv[3], "GPU") == 0) {
+    LOG(ERROR) << "Using GPU";
+    uint device_id = 0;
+    if (argc >= 5 && strcmp(argv[3], "GPU") == 0) {
+      device_id = atoi(argv[4]);
+    }
+    LOG(ERROR) << "Using Device_id=" << device_id;
+    Caffe::SetDevice(device_id);
+    Caffe::set_mode(Caffe::GPU);
+  } else {
+    LOG(ERROR) << "Using CPU";
+    Caffe::set_mode(Caffe::CPU);
+  }
+
+  Caffe::set_phase(Caffe::TRAIN);
   NetParameter net_param;
   ReadProtoFromTextFile(argv[1],
       &net_param);
@@ -40,24 +64,29 @@ int main(int argc, char** argv) {
   vector<vector<Blob<float>*> >& bottom_vecs = caffe_net.bottom_vecs();
   vector<vector<Blob<float>*> >& top_vecs = caffe_net.top_vecs();
   LOG(ERROR) << "*** Benchmark begins ***";
+  clock_t forward_start = clock();
   for (int i = 0; i < layers.size(); ++i) {
     const string& layername = layers[i]->layer_param().name();
     clock_t start = clock();
-    for (int j = 0; j < repeat; ++j) {
+    for (int j = 0; j < total_iter; ++j) {
       layers[i]->Forward(bottom_vecs[i], &top_vecs[i]);
     }
     LOG(ERROR) << layername << "\tforward: "
         << float(clock() - start) / CLOCKS_PER_SEC << " seconds.";
   }
+  LOG(ERROR) << "Forward pass: " << float(clock() - forward_start) / CLOCKS_PER_SEC << " seconds.";
+  clock_t backward_start = clock();
   for (int i = layers.size() - 1; i >= 0; --i) {
     const string& layername = layers[i]->layer_param().name();
     clock_t start = clock();
-    for (int j = 0; j < repeat; ++j) {
+    for (int j = 0; j < total_iter; ++j) {
       layers[i]->Backward(top_vecs[i], true, &bottom_vecs[i]);
     }
     LOG(ERROR) << layername << "\tbackward: "
         << float(clock() - start) / CLOCKS_PER_SEC << " seconds.";
   }
+  LOG(ERROR) << "Backward pass: " << float(clock() - backward_start) / CLOCKS_PER_SEC << " seconds.";
+  LOG(ERROR) << "Total Time: " << float(clock() - forward_start) / CLOCKS_PER_SEC << " seconds.";
   LOG(ERROR) << "*** Benchmark ends ***";
   return 0;
 }