[pycaffe] expose SGDSolver.step

BVLC · Jan 3, 2015 · 10c2364 · 10c2364
1 parent 55d7ffd
commit 10c2364
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
@@ -197,7 +197,8 @@ BOOST_PYTHON_MODULE(_caffe) {
       .add_property("test_nets", &PySGDSolver::test_nets)
       .add_property("iter",      &PySGDSolver::iter)
       .def("solve",              &PySGDSolver::Solve)
-      .def("solve",              &PySGDSolver::SolveResume);
+      .def("solve",              &PySGDSolver::SolveResume)
+      .def("step",               &PySGDSolver::Step);
 
   bp::class_<vector<shared_ptr<PyNet> > >("NetVec")
       .def(bp::vector_indexing_suite<vector<shared_ptr<PyNet> >, true>());

diff --git a/python/caffe/_caffe.hpp b/python/caffe/_caffe.hpp
@@ -181,6 +181,7 @@ class PySGDSolver {
   vector<shared_ptr<PyNet> > test_nets() { return test_nets_; }
   int iter() { return solver_->iter(); }
   void Solve() { return solver_->Solve(); }
+  void Step(int iters) { solver_->Step(iters); }
   void SolveResume(const string& resume_file);
 
  protected: