collect Net inputs from Input layers

Restore the list of net inputs for compatibility with the pycaffe and matcaffe interfaces and downstream C++.
chengyangfu · Feb 25, 2016 · 0d9a78f · 0d9a78f
1 parent 51f79a8
commit 0d9a78f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/include/caffe/net.hpp b/include/caffe/net.hpp
@@ -185,11 +185,18 @@ class Net {
   inline const vector<string>& param_display_names() const {
     return param_display_names_;
   }
-  /// @brief output blob number
+  /// @brief Input and output blob numbers
+  inline int num_inputs() const { return net_input_blobs_.size(); }
   inline int num_outputs() const { return net_output_blobs_.size(); }
+  inline const vector<Blob<Dtype>*>& input_blobs() const {
+    return net_input_blobs_;
+  }
   inline const vector<Blob<Dtype>*>& output_blobs() const {
     return net_output_blobs_;
   }
+  inline const vector<int>& input_blob_indices() const {
+    return net_input_blob_indices_;
+  }
   inline const vector<int>& output_blob_indices() const {
     return net_output_blob_indices_;
   }
@@ -263,8 +270,10 @@ class Net {
   vector<string> param_display_names_;
   vector<pair<int, int> > param_layer_indices_;
   map<string, int> param_names_index_;
-  /// blob indices for the output of the net
+  /// blob indices for the input and the output of the net
+  vector<int> net_input_blob_indices_;
   vector<int> net_output_blob_indices_;
+  vector<Blob<Dtype>*> net_input_blobs_;
   vector<Blob<Dtype>*> net_output_blobs_;
   /// The parameters in the network.
   vector<shared_ptr<Blob<Dtype> > > params_;

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
@@ -103,6 +103,12 @@ void Net<Dtype>::Init(const NetParameter& in_param) {
     int num_top = layer_param.top_size();
     for (int top_id = 0; top_id < num_top; ++top_id) {
       AppendTop(param, layer_id, top_id, &available_blobs, &blob_name_to_idx);
+      // Collect Input layer tops as Net inputs.
+      if (layer_param.type() == "Input") {
+        const int blob_id = blobs_.size() - 1;
+        net_input_blob_indices_.push_back(blob_id);
+        net_input_blobs_.push_back(blobs_[blob_id].get());
+      }
     }
     // If the layer specifies that AutoTopBlobs() -> true and the LayerParameter
     // specified fewer than the required number (as specified by