add split layer insertion tests; move split insertion code to util file

3 files changed, 401 insertions, 67 deletions

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
index 3c4148fe..fbb109fe 100644
--- a/src/caffe/net.cpp
+++ b/src/caffe/net.cpp
@@ -10,6 +10,7 @@
 #include "caffe/layer.hpp"
 #include "caffe/net.hpp"
 #include "caffe/util/io.hpp"
+#include "caffe/util/insert_splits.hpp"
 
 using std::pair;
 using std::map;
@@ -33,7 +34,7 @@ template <typename Dtype>
 void Net<Dtype>::Init(const NetParameter& in_param) {
   // Create a copy of in_param with splits added where necessary.
   NetParameter param;
-  AddSplits(in_param, &param);
+  insert_splits(in_param, &param);
   // Basically, build all the layers and set up its connections.
   name_ = param.name();
   map<string, int> blob_name_to_idx;
@@ -158,72 +159,6 @@ void Net<Dtype>::Init(const NetParameter& in_param) {
 
 
 template <typename Dtype>
-void Net<Dtype>::AddSplits(const NetParameter& param,
-    NetParameter* param_split) {
-  // Initialize by copying from the input NetParameter.
-  param_split->CopyFrom(param);
-  param_split->clear_layers();
-  map<string, int> blob_name_to_bottom_count;
-  map<string, int> blob_name_to_bottom_split_idx;
-  // Determine for each top blob the number of times it's used as a bottom blob.
-  for (int i = 0; i < param.layers_size(); ++i) {
-    const LayerConnection& layer_connection = param.layers(i);
-    for (int j = 0; j < layer_connection.bottom_size(); ++j) {
-      const string& blob_name = layer_connection.bottom(j);
-      blob_name_to_bottom_count[blob_name]++;
-    }
-    for (int j = 0; j < layer_connection.top_size(); ++j) {
-      const string& blob_name = layer_connection.top(j);
-      blob_name_to_bottom_count[blob_name] = 0;
-      blob_name_to_bottom_split_idx[blob_name] = 0;
-    }
-  }
-  for (int i = 0; i < param.layers_size(); ++i) {
-    LayerConnection* layer_connection = param_split->add_layers();
-    layer_connection->CopyFrom(param.layers(i));
-    // Replace any shared bottom blobs with split layer outputs.
-    for (int j = 0; j < layer_connection->bottom_size(); ++j) {
-      const string& blob_name = layer_connection->bottom(j);
-      const int split_count = blob_name_to_bottom_count[blob_name];
-      if (split_count > 1) {
-        const int suffix_max_length = 16;
-        char split_suffix[suffix_max_length];
-        const int suffix_length = snprintf(split_suffix, suffix_max_length,
-            "_split_%d", blob_name_to_bottom_split_idx[blob_name]++);
-        CHECK_LT(suffix_length, suffix_max_length);
-        const string& split_blob_name = blob_name + split_suffix;
-        layer_connection->set_bottom(j, split_blob_name);
-      }
-    }
-    // Create split blob for any top blobs used by other layers as bottom
-    // blobs more than once.
-    for (int j = 0; j < layer_connection->top_size(); ++j) {
-      const string& blob_name = layer_connection->top(j);
-      const int split_count = blob_name_to_bottom_count[blob_name];
-      if (split_count > 1) {
-        LayerConnection* split_layer_connection = param_split->add_layers();
-        split_layer_connection->add_bottom(blob_name);
-        LayerParameter* split_layer_param =
-            split_layer_connection->mutable_layer();
-        split_layer_param->set_name(blob_name + "_split");
-        split_layer_param->set_type("split");
-        vector<Blob<Dtype>*> split_top_blobs(split_count);
-        for (int k = 0; k < split_count; ++k) {
-          const int suffix_max_length = 16;
-          char split_suffix[suffix_max_length];
-          const int suffix_length = snprintf(split_suffix, suffix_max_length,
-              "_split_%d", k);
-          CHECK_LT(suffix_length, suffix_max_length);
-          const string& split_blob_name = blob_name + split_suffix;
-          split_layer_connection->add_top(split_blob_name);
-        }
-      }
-    }
-  }
-}
-
-
-template <typename Dtype>
 void Net<Dtype>::GetLearningRateAndWeightDecay() {
   LOG(INFO) << "Collecting Learning Rate and Weight Decay.";
   for (int i = 0; i < layers_.size(); ++i) {
diff --git a/src/caffe/test/test_split_layer.cpp b/src/caffe/test/test_split_layer.cpp
index 9a3a2d37..aefa50c5 100644
--- a/src/caffe/test/test_split_layer.cpp
+++ b/src/caffe/test/test_split_layer.cpp
@@ -2,6 +2,7 @@
 
 #include <cstring>
 #include <cuda_runtime.h>
+#include <google/protobuf/text_format.h>
 
 #include "gtest/gtest.h"
 #include "caffe/blob.hpp"
@@ -9,6 +10,7 @@
 #include "caffe/filler.hpp"
 #include "caffe/vision_layers.hpp"
 #include "caffe/test/test_gradient_check_util.hpp"
+#include "caffe/util/insert_splits.hpp"
 
 #include "caffe/test/test_caffe_main.hpp"
 
@@ -106,4 +108,304 @@ TYPED_TEST(SplitLayerTest, TestGPUGradient) {
       this->blob_top_vec_);
 }
 
+
+template <typename Dtype>
+class SplitLayerInsertionTest : public ::testing::Test {
+ protected:
+ SplitLayerInsertionTest() { };
+  void RunInsertionTest(
+      const string& input_param_string, const string& output_param_string) {
+    NetParameter input_param;
+    CHECK(google::protobuf::TextFormat::ParseFromString(
+        input_param_string, &input_param));
+    NetParameter expected_output_param;
+    CHECK(google::protobuf::TextFormat::ParseFromString(
+        output_param_string, &expected_output_param));
+    NetParameter actual_output_param;
+    insert_splits(input_param, &actual_output_param);
+    CHECK_EQ(expected_output_param.DebugString(),
+        actual_output_param.DebugString());
+    EXPECT_EQ(expected_output_param.DebugString(),
+        actual_output_param.DebugString());
+  }
+};
+
+typedef ::testing::Types<float> InsertionDtypes;
+TYPED_TEST_CASE(SplitLayerInsertionTest, InsertionDtypes);
+
+TYPED_TEST(SplitLayerInsertionTest, TestNoInsertion1) {
+  const string& input_proto =
+      "name: \"TestNetwork\"\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data\"\n"
+      "    type: \"data\"\n"
+      "  }\n"
+      "  top: \"data\"\n"
+      "  top: \"label\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"softmax_with_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod\"\n"
+      "  bottom: \"label\"\n"
+      "}\n";
+  this->RunInsertionTest(input_proto, input_proto);
+}
+
+TYPED_TEST(SplitLayerInsertionTest, TestNoInsertion2) {
+  const string& input_proto =
+      "name: \"TestNetwork\"\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data\"\n"
+      "    type: \"data\"\n"
+      "  }\n"
+      "  top: \"data\"\n"
+      "  top: \"label\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data_split\"\n"
+      "    type: \"split\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"data_split_0\"\n"
+      "  top: \"data_split_1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod1\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_0\"\n"
+      "  top: \"innerprod1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_1\"\n"
+      "  top: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod1\"\n"
+      "  bottom: \"innerprod2\"\n"
+      "}\n";
+  this->RunInsertionTest(input_proto, input_proto);
+}
+
+TYPED_TEST(SplitLayerInsertionTest, TestInsertion) {
+  const string& input_proto =
+      "name: \"TestNetwork\"\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data\"\n"
+      "    type: \"data\"\n"
+      "  }\n"
+      "  top: \"data\"\n"
+      "  top: \"label\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod1\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod3\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod3\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod1\"\n"
+      "  bottom: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod2\"\n"
+      "  bottom: \"innerprod3\"\n"
+      "}\n";
+  const string& expected_output_proto =
+      "name: \"TestNetwork\"\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data\"\n"
+      "    type: \"data\"\n"
+      "  }\n"
+      "  top: \"data\"\n"
+      "  top: \"label\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data_split\"\n"
+      "    type: \"split\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"data_split_0\"\n"
+      "  top: \"data_split_1\"\n"
+      "  top: \"data_split_2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod1\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_0\"\n"
+      "  top: \"innerprod1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_1\"\n"
+      "  top: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2_split\"\n"
+      "    type: \"split\"\n"
+      "  }\n"
+      "  bottom: \"innerprod2\"\n"
+      "  top: \"innerprod2_split_0\"\n"
+      "  top: \"innerprod2_split_1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod3\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_2\"\n"
+      "  top: \"innerprod3\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod1\"\n"
+      "  bottom: \"innerprod2_split_0\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod2_split_1\"\n"
+      "  bottom: \"innerprod3\"\n"
+      "}\n";
+  this->RunInsertionTest(input_proto, expected_output_proto);
+}
+
+TYPED_TEST(SplitLayerInsertionTest, TestInputInsertion) {
+  const string& input_proto =
+      "name: \"TestNetwork\"\n"
+      "input: \"data\"\n"
+      "input_dim: 10\n"
+      "input_dim: 3\n"
+      "input_dim: 227\n"
+      "input_dim: 227\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod1\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod1\"\n"
+      "  bottom: \"innerprod2\"\n"
+      "}\n";
+  const string& expected_output_proto =
+      "name: \"TestNetwork\"\n"
+      "input: \"data\"\n"
+      "input_dim: 10\n"
+      "input_dim: 3\n"
+      "input_dim: 227\n"
+      "input_dim: 227\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"data_split\"\n"
+      "    type: \"split\"\n"
+      "  }\n"
+      "  bottom: \"data\"\n"
+      "  top: \"data_split_0\"\n"
+      "  top: \"data_split_1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod1\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_0\"\n"
+      "  top: \"innerprod1\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"innerprod2\"\n"
+      "    type: \"inner_product\"\n"
+      "  }\n"
+      "  bottom: \"data_split_1\"\n"
+      "  top: \"innerprod2\"\n"
+      "}\n"
+      "layers: {\n"
+      "  layer {\n"
+      "    name: \"loss\"\n"
+      "    type: \"euclidean_loss\"\n"
+      "  }\n"
+      "  bottom: \"innerprod1\"\n"
+      "  bottom: \"innerprod2\"\n"
+      "}\n";
+  this->RunInsertionTest(input_proto, expected_output_proto);
+}
+
 }
diff --git a/src/caffe/util/insert_splits.cpp b/src/caffe/util/insert_splits.cpp
new file mode 100644
index 00000000..df8bf5da
--- /dev/null
+++ b/src/caffe/util/insert_splits.cpp
@@ -0,0 +1,97 @@
+// Copyright 2014 Jeff Donahue
+
+#include <fstream>
+#include <map>
+#include <string>
+
+#include "caffe/common.hpp"
+#include "caffe/util/insert_splits.hpp"
+
+using std::map;
+
+namespace caffe {
+
+void insert_splits(const NetParameter& param, NetParameter* param_split) {
+  // Initialize by copying from the input NetParameter.
+  param_split->CopyFrom(param);
+  param_split->clear_layers();
+  map<string, int> blob_name_to_bottom_count;
+  map<string, int> blob_name_to_bottom_split_idx;
+  // Determine for each top blob (including input blobs) the number of times
+  // it's used as a bottom blob.
+  for (int i = 0; i < param.input_size(); ++i) {
+    const string& blob_name = param.input(i);
+    blob_name_to_bottom_count[blob_name] = 0;
+    blob_name_to_bottom_split_idx[blob_name] = 0;
+  }
+  for (int i = 0; i < param.layers_size(); ++i) {
+    const LayerConnection& layer_connection = param.layers(i);
+    for (int j = 0; j < layer_connection.bottom_size(); ++j) {
+      const string& blob_name = layer_connection.bottom(j);
+      blob_name_to_bottom_count[blob_name]++;
+    }
+    for (int j = 0; j < layer_connection.top_size(); ++j) {
+      const string& blob_name = layer_connection.top(j);
+      blob_name_to_bottom_count[blob_name] = 0;
+      blob_name_to_bottom_split_idx[blob_name] = 0;
+    }
+  }
+  for (int i = 0; i < param.input_size(); ++i) {
+    const string& blob_name = param.input(i);
+    const int split_count = blob_name_to_bottom_count[blob_name];
+    if (split_count > 1) {
+      LayerConnection* split_layer_connection = param_split->add_layers();
+      configure_split_layer(blob_name, split_count, split_layer_connection);
+    }
+  }
+  for (int i = 0; i < param.layers_size(); ++i) {
+    LayerConnection* layer_connection = param_split->add_layers();
+    layer_connection->CopyFrom(param.layers(i));
+    // Replace any shared bottom blobs with split layer outputs.
+    for (int j = 0; j < layer_connection->bottom_size(); ++j) {
+      const string& blob_name = layer_connection->bottom(j);
+      const int split_count = blob_name_to_bottom_count[blob_name];
+      if (split_count > 1) {
+        const int suffix_max_length = 16;
+        char split_suffix[suffix_max_length];
+        const int suffix_length = snprintf(split_suffix, suffix_max_length,
+            "_split_%d", blob_name_to_bottom_split_idx[blob_name]++);
+        CHECK_LT(suffix_length, suffix_max_length);
+        const string& split_blob_name = blob_name + split_suffix;
+        layer_connection->set_bottom(j, split_blob_name);
+      }
+    }
+    // Create split blob for any top blobs used by other layers as bottom
+    // blobs more than once.
+    for (int j = 0; j < layer_connection->top_size(); ++j) {
+      const string& blob_name = layer_connection->top(j);
+      const int split_count = blob_name_to_bottom_count[blob_name];
+      if (split_count > 1) {
+        LayerConnection* split_layer_connection = param_split->add_layers();
+        split_layer_connection->add_bottom(blob_name);
+        configure_split_layer(blob_name, split_count, split_layer_connection);
+      }
+    }
+  }
+}
+
+void configure_split_layer(const string& blob_name,
+    const int split_count, LayerConnection* split_layer_connection) {
+  split_layer_connection->Clear();
+  split_layer_connection->add_bottom(blob_name);
+  LayerParameter* split_layer_param =
+      split_layer_connection->mutable_layer();
+  split_layer_param->set_name(blob_name + "_split");
+  split_layer_param->set_type("split");
+  for (int k = 0; k < split_count; ++k) {
+    const int suffix_max_length = 16;
+    char split_suffix[suffix_max_length];
+    const int suffix_length = snprintf(split_suffix, suffix_max_length,
+        "_split_%d", k);
+    CHECK_LT(suffix_length, suffix_max_length);
+    const string& split_blob_name = blob_name + split_suffix;
+    split_layer_connection->add_top(split_blob_name);
+  }
+}
+
+}  // namespace caffe