Merge pull request BVLC#4348 from nitnelave/python/clear_params

Add clear_param_diffs to the python net interface
andpol5 · Jul 11, 2016 · a1a9d30 · a1a9d30
2 parents 4af9062 + 892c78d
commit a1a9d30
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
@@ -299,6 +299,7 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def("_forward", &Net<Dtype>::ForwardFromTo)
     .def("_backward", &Net<Dtype>::BackwardFromTo)
     .def("reshape", &Net<Dtype>::Reshape)
+    .def("clear_param_diffs", &Net<Dtype>::ClearParamDiffs)
     // The cast is to select a particular overload.
     .def("copy_from", static_cast<void (Net<Dtype>::*)(const string)>(
         &Net<Dtype>::CopyTrainedLayersFrom))

diff --git a/python/caffe/test/test_net.py b/python/caffe/test/test_net.py
@@ -63,6 +63,17 @@ def test_forward_backward(self):
         self.net.forward()
         self.net.backward()
 
+    def test_clear_param_diffs(self):
+        # Run a forward/backward step to have non-zero diffs
+        self.net.forward()
+        self.net.backward()
+        diff = self.net.params["conv"][0].diff
+        # Check that we have non-zero diffs
+        self.assertTrue(diff.max() > 0)
+        self.net.clear_param_diffs()
+        # Check that the diffs are now 0
+        self.assertTrue((diff == 0).all())
+
     def test_inputs_outputs(self):
         self.assertEqual(self.net.inputs, [])
         self.assertEqual(self.net.outputs, ['loss'])