pycaffe: store a shared_ptr<CaffeNet> in SGDSolver

Doing this, rather than constructing the CaffeNet wrapper every time,
will allow the wrapper to hold references that last at least as long as
SGDSolver (which will be necessary to ensure that data used by
MemoryDataLayer doesn't get freed).

1 files changed, 5 insertions, 1 deletions

diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
index 1a44debf..853ddbe1 100644
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
@@ -301,9 +301,12 @@ class CaffeSGDSolver {
     // exception if param_file can't be opened
     CheckFile(param_file);
     solver_.reset(new SGDSolver<float>(param_file));
+    // we need to explicitly store the net wrapper, rather than constructing
+    // it on the fly, so that it can hold references to Python objects
+    net_.reset(new CaffeNet(solver_->net()));
   }
 
-  CaffeNet net() { return CaffeNet(solver_->net()); }
+  shared_ptr<CaffeNet> net() { return net_; }
   void Solve() { return solver_->Solve(); }
   void SolveResume(const string& resume_file) {
     CheckFile(resume_file);
@@ -311,6 +314,7 @@ class CaffeSGDSolver {
   }
 
  protected:
+  shared_ptr<CaffeNet> net_;
   shared_ptr<SGDSolver<float> > solver_;
 };