specify NetParameters directly in the SolverParameter

BVLC · May 10, 2014 · 65ef9ff · 65ef9ff
1 parent 6fcf5ed
commit 65ef9ff
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto
@@ -57,8 +57,16 @@ message NetParameter {
 }
 
 message SolverParameter {
-  optional string train_net = 1; // The proto file for the training net.
-  optional string test_net = 2; // The proto file for the testing net.
+  // {train,test}_net specify a path to a file containing the {train,test} net
+  // parameters; {train,test}_net_param specify the net parameters directly
+  // inside the SolverParameter.
+  //
+  // If {train,test}_net is specified, {train,test}_net_param should not be,
+  // and vice versa.
+  optional string train_net = 1; // The proto filename for the train net.
+  optional string test_net = 2; // The proto filename for the test net.
+  optional NetParameter train_net_param = 21; // Full params for the train net.
+  optional NetParameter test_net_param = 22; // Full params for the test net.
   // The number of iterations for each testing phase.
   optional int32 test_iter = 3 [default = 0];
   // The number of iterations between two testing phases.

diff --git a/src/caffe/solver.cpp b/src/caffe/solver.cpp
@@ -38,11 +38,29 @@ void Solver<Dtype>::Init(const SolverParameter& param) {
     Caffe::set_random_seed(param_.random_seed());
   }
   // Scaffolding code
-  LOG(INFO) << "Creating training net.";
-  net_.reset(new Net<Dtype>(param_.train_net()));
-  if (param_.has_test_net()) {
-    LOG(INFO) << "Creating testing net.";
+  if (param_.has_train_net_param()) {
+    CHECK(!param_.has_train_net()) << "Either train_net_param or train_net may "
+                                   << "be specified, but not both.";
+    LOG(INFO) << "Creating training net specified in SolverParameter.";
+    net_.reset(new Net<Dtype>(param_.train_net_param()));
+  } else {
+    LOG(INFO) << "Creating training net from file: " << param_.train_net();
+    net_.reset(new Net<Dtype>(param_.train_net()));
+  }
+  bool has_test_net = false;
+  NetParameter test_net_param;
+  if (param_.has_test_net_param()) {
+    CHECK(!param_.has_test_net()) << "Either test_net_param or test_net may be "
+                                  << "specified, but not both.";
+    LOG(INFO) << "Creating testing net specified in SolverParameter.";
+    test_net_.reset(new Net<Dtype>(param_.test_net_param()));
+    has_test_net = true;
+  } else if (param_.has_test_net()) {
+    LOG(INFO) << "Creating testing net from file: " << param_.test_net();
     test_net_.reset(new Net<Dtype>(param_.test_net()));
+    has_test_net = true;
+  }
+  if (has_test_net) {
     CHECK_GT(param_.test_iter(), 0);
     CHECK_GT(param_.test_interval(), 0);
   }