Exposing blob loss weight to python

BVLC · Jul 23, 2015 · b3badcd · b3badcd
1 parent ad8d129
commit b3badcd
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
@@ -211,6 +211,8 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def("copy_from", static_cast<void (Net<Dtype>::*)(const string)>(
         &Net<Dtype>::CopyTrainedLayersFrom))
     .def("share_with", &Net<Dtype>::ShareTrainedLayersWith)
+    .add_property("_blob_loss_weights", bp::make_function(
+        &Net<Dtype>::blob_loss_weights, bp::return_internal_reference<>()))
     .add_property("_blobs", bp::make_function(&Net<Dtype>::blobs,
         bp::return_internal_reference<>()))
     .add_property("layers", bp::make_function(&Net<Dtype>::layers,
@@ -288,6 +290,8 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def(bp::vector_indexing_suite<vector<string> >());
   bp::class_<vector<int> >("IntVec")
     .def(bp::vector_indexing_suite<vector<int> >());
+  bp::class_<vector<Dtype> >("DtypeVec")
+    .def(bp::vector_indexing_suite<vector<Dtype> >());
   bp::class_<vector<shared_ptr<Net<Dtype> > > >("NetVec")
     .def(bp::vector_indexing_suite<vector<shared_ptr<Net<Dtype> > >, true>());
   bp::class_<vector<bool> >("BoolVec")

diff --git a/python/caffe/pycaffe.py b/python/caffe/pycaffe.py
@@ -27,6 +27,15 @@ def _Net_blobs(self):
     return OrderedDict(zip(self._blob_names, self._blobs))
 
 
+@property
+def _Net_blob_loss_weights(self):
+    """
+    An OrderedDict (bottom to top, i.e., input to output) of network
+    blob loss weights indexed by name
+    """
+    return OrderedDict(zip(self._blob_names, self._blob_loss_weights))
+
+
 @property
 def _Net_params(self):
     """
@@ -270,6 +279,7 @@ def _Net_batch(self, blobs):
 
 # Attach methods to Net.
 Net.blobs = _Net_blobs
+Net.blob_loss_weights = _Net_blob_loss_weights
 Net.params = _Net_params
 Net.forward = _Net_forward
 Net.backward = _Net_backward