Add 'snapshot_after_train' to SolverParameter to override the final

snapshot.

diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto
index d2f61e5..4e77863 100644 (file)
--- a/src/caffe/proto/caffe.proto
+++ b/src/caffe/proto/caffe.proto
@@ -140,6 +140,9 @@ message SolverParameter {
   // If true, print information about the state of the net that may help with
   // debugging learning problems.
   optional bool debug_info = 23 [default = false];
+
+  // If false, don't save a snapshot after training finishes.
+  optional bool snapshot_after_train = 28 [default = true];
 }
 
 // A message that stores the solver snapshots

diff --git a/src/caffe/solver.cpp b/src/caffe/solver.cpp
index 11795f0..6049ddf 100644 (file)
--- a/src/caffe/solver.cpp
+++ b/src/caffe/solver.cpp
@@ -200,8 +200,9 @@ void Solver<Dtype>::Solve(const char* resume_file) {
     ComputeUpdateValue();
     net_->Update();
   }
-  // Always save a snapshot after optimization.
-  Snapshot();
+  // Always save a snapshot after optimization, unless overridden by setting
+  // snapshot_after_train := false.
+  if (param_.snapshot_after_train()) { Snapshot(); }
   // After the optimization is done, run an additional train and test pass to
   // display the train and test loss/outputs if appropriate (based on the
   // display and test_interval settings, respectively).  Unlike in the rest of