Exposing HDF5 saving and loading to python

[philkr]

Exposes a new function 'Net.saveHDF5' to python. This allows caffe to store model weights in the hdf5 format (which can easily be used in other libraries).

[ajtulloch]

Would you considering updating this to also add loadHDF5 as well? If the point of this is to save models with tensors > 2GB, then it would be nice to be able to load them in PyCaffe as well.

Here's a patch that I'd suggest modifying yours to:

diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
@@ -101,6 +101,14 @@
   WriteProtoToBinaryFile(net_param, filename.c_str());
 }

+void Net_SaveHDF5(const Net<Dtype>& net, string filename) {
+  net.ToHDF5(filename.c_str(), false);
+}
+
+void Net_LoadHDF5(Net<Dtype>* net, string filename) {
+  net->CopyTrainedLayersFromHDF5(filename.c_str());
+}
+
 void Net_SetInputArrays(Net<Dtype>* net, bp::object data_obj,
     bp::object labels_obj) {
   // check that this network has an input MemoryDataLayer
@@ -254,6 +262,8 @@
         bp::return_value_policy<bp::copy_const_reference>()))
     .def("_set_input_arrays", &Net_SetInputArrays,
         bp::with_custodian_and_ward<1, 2, bp::with_custodian_and_ward<1, 3> >())
+    .def("load_hdf5", &Net_LoadHDF5)
+    .def("save_hdf5", &Net_SaveHDF5)
     .def("save", &Net_Save);

   bp::class_<Blob<Dtype>, shared_ptr<Blob<Dtype> >, boost::noncopyable>(
diff --git a/python/caffe/test/test_net.py b/python/caffe/test/test_net.py
--- a/python/caffe/test/test_net.py
+++ b/python/caffe/test/test_net.py
@@ -79,3 +79,17 @@
             for i in range(len(self.net.params[name])):
                 self.assertEqual(abs(self.net.params[name][i].data
                     - net2.params[name][i].data).sum(), 0)
+
+    def test_save_hdf5(self):
+        f = tempfile.NamedTemporaryFile(mode='w+', delete=False)
+        f.close()
+        self.net.save_hdf5(f.name)
+        net_file = simple_net_file(self.num_output)
+        net2 = caffe.Net(net_file, caffe.TRAIN)
+        net2.load_hdf5(f.name)
+        os.remove(net_file)
+        os.remove(f.name)
+        for name in self.net.params:
+            for i in range(len(self.net.params[name])):
+                self.assertEqual(abs(self.net.params[name][i].data
+                    - net2.params[name][i].data).sum(), 0)

[philkr]

thank you @ajtulloch , I just added the load_hdf5 function.

[ajtulloch]

LGTM. @longjon, @shelhamer?

[ajtulloch]

You're using camel case in the _caffe.cpp file (saveHDF5, loadHDF5), but snake case here. I think for consistency you should use save_hdf5, load_hdf5 in the _caffe.cpp file. It's Python PEP-8 style to use snake_case for member functions.

[ajtulloch]

Could you fix the casing issue? That will fix travis.

[shelhamer]

@philkr please squash your commits for merge. We should expose hdf5 load/save to resolve #4192 and for the purpose of interoperability as you mentioned.

This is straightforward, but @longjon could also double-check the diff.

[shelhamer]

@philkr I think you meant Exposing save_hdf5 and load_hdf5 to python for the commit message.

[philkr]

@shelhamer yes I did, changed it now.

[shelhamer]

Thanks Philipp for exposing hdf5 net serialization to pycaffe.

[ajtulloch]

Thanks @philkr!