Expose tops and bottoms associated with layers in python, fix #2829

include/caffe/net.hpp

@@ -124,13 +124,19 @@ class Net {
   inline const vector<vector<Blob<Dtype>*> >& bottom_vecs() const {
     return bottom_vecs_;
   }
+  inline const vector<vector<int> >& bottom_id_vecs() const {
+    return bottom_id_vecs_;
+  }
   /**
    * @brief returns the top vecs for each layer -- usually you won't
    *        need this unless you do per-layer checks such as gradients.
    */
   inline const vector<vector<Blob<Dtype>*> >& top_vecs() const {
     return top_vecs_;
   }
+  inline const vector<vector<int> >& top_id_vecs() const {
+    return top_id_vecs_;
+  }
   inline const vector<vector<bool> >& bottom_need_backward() const {
     return bottom_need_backward_;
   }

python/caffe/_caffe.cpp

@@ -131,6 +131,16 @@ void Net_SetInputArrays(Net<Dtype>* net, bp::object data_obj,
       PyArray_DIMS(data_arr)[0]);
 }
 
+const vector<int>& Net_BottomsForLayer(const Net<Dtype>& net,
+                                                int layer_idx) {
+  return net.bottom_id_vecs()[layer_idx];
+}
+
+const vector<int>& Net_TopsForLayer(const Net<Dtype>& net,
+                                             int layer_idx) {
+  return net.top_id_vecs()[layer_idx];
+}
+
 Solver<Dtype>* GetSolverFromFile(const string& filename) {
   SolverParameter param;
   ReadProtoFromTextFileOrDie(filename, &param);
@@ -226,7 +236,11 @@ BOOST_PYTHON_MODULE(_caffe) {
         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("save", &Net_Save);
+    .def("save", &Net_Save)
+    .def("_bottoms_for_layer", &Net_BottomsForLayer,
+        bp::return_value_policy<bp::copy_const_reference>())
+    .def("_tops_for_layer", &Net_TopsForLayer,
+        bp::return_value_policy<bp::copy_const_reference>());
 
   bp::class_<Blob<Dtype>, shared_ptr<Blob<Dtype> >, boost::noncopyable>(
     "Blob", bp::no_init)

python/caffe/pycaffe.py

@@ -77,7 +77,9 @@ def _Net_forward(self, blobs=None, start=None, end=None, **kwargs):
 
     if end is not None:
         end_ind = list(self._layer_names).index(end)
-        outputs = set([end] + blobs)
+        outputs = set(blobs)
+        outputs.update([self._blob_names[idx]
+                        for idx in self._tops_for_layer(end_ind)])
     else:
         end_ind = len(self.layers) - 1
         outputs = set(self.outputs + blobs)
@@ -125,7 +127,9 @@ def _Net_backward(self, diffs=None, start=None, end=None, **kwargs):
 
     if end is not None:
         end_ind = list(self._layer_names).index(end)
-        outputs = set([end] + diffs)
+        outputs = set(diffs)
+        outputs.update([self._blob_names[idx]
+                        for idx in self._bottoms_for_layer(end_ind)])
     else:
         end_ind = 0
         outputs = set(self.inputs + diffs)
@@ -268,6 +272,18 @@ def _Net_batch(self, blobs):
                                                  padding])
         yield padded_batch
 
+def _Net_bottoms_for_layer(self, layer_name):
+    layer_idx = list(self._layer_names).index(layer_name);
+    blob_idx = self._bottoms_for_layer(layer_idx);
+    return OrderedDict([(self._blob_names[idx], self._blobs[idx])
+                        for idx in blob_idx]);
+
+def _Net_tops_for_layer(self, layer_name):
+    layer_idx = list(self._layer_names).index(layer_name);
+    blob_idx = self._tops_for_layer(layer_idx);
+    return OrderedDict([(self._blob_names[idx], self._blobs[idx])
+                        for idx in blob_idx]);
+
 # Attach methods to Net.
 Net.blobs = _Net_blobs
 Net.params = _Net_params
@@ -279,3 +295,5 @@ def _Net_batch(self, blobs):
 Net._batch = _Net_batch
 Net.inputs = _Net_inputs
 Net.outputs = _Net_outputs
+Net.bottoms_for_layer = _Net_bottoms_for_layer
+Net.tops_for_layer = _Net_tops_for_layer