diff --git a/include/caffe/net.hpp b/include/caffe/net.hpp index 5665df1edf2..c8b90661015 100644 --- a/include/caffe/net.hpp +++ b/include/caffe/net.hpp @@ -124,6 +124,9 @@ class Net { inline const vector*> >& bottom_vecs() const { return bottom_vecs_; } + inline const vector >& 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. @@ -131,6 +134,9 @@ class Net { inline const vector*> >& top_vecs() const { return top_vecs_; } + inline const vector >& top_id_vecs() const { + return top_id_vecs_; + } inline const vector >& bottom_need_backward() const { return bottom_need_backward_; } diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp index dff7f627016..86f9ecf6e9f 100644 --- a/python/caffe/_caffe.cpp +++ b/python/caffe/_caffe.cpp @@ -131,6 +131,16 @@ void Net_SetInputArrays(Net* net, bp::object data_obj, PyArray_DIMS(data_arr)[0]); } +const vector& Net_BottomsForLayer(const Net& net, + int layer_idx) { + return net.bottom_id_vecs()[layer_idx]; +} + +const vector& Net_TopsForLayer(const Net& net, + int layer_idx) { + return net.top_id_vecs()[layer_idx]; +} + Solver* GetSolverFromFile(const string& filename) { SolverParameter param; ReadProtoFromTextFileOrDie(filename, ¶m); @@ -226,7 +236,11 @@ BOOST_PYTHON_MODULE(_caffe) { bp::return_value_policy())) .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()) + .def("_tops_for_layer", &Net_TopsForLayer, + bp::return_value_policy()); bp::class_, shared_ptr >, boost::noncopyable>( "Blob", bp::no_init) diff --git a/python/caffe/pycaffe.py b/python/caffe/pycaffe.py index e8a676a26d2..15d13eaf5a3 100644 --- a/python/caffe/pycaffe.py +++ b/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