Define split layer (merge trick)

Originally authored by Jeff Donahue, but this commit was needed for historical reasons.
noruya · Feb 26, 2014 · 527cfab · 527cfab
1 parent a8c7407
commit 527cfab
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/include/caffe/vision_layers.hpp b/include/caffe/vision_layers.hpp
@@ -125,6 +125,27 @@ class DropoutLayer : public NeuronLayer<Dtype> {
 };
 
 
+template <typename Dtype>
+class SplitLayer : public Layer<Dtype> {
+ public:
+  explicit SplitLayer(const LayerParameter& param)
+      : Layer<Dtype>(param) {}
+  virtual void SetUp(const vector<Blob<Dtype>*>& bottom,
+      vector<Blob<Dtype>*>* top);
+
+ protected:
+  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
+      vector<Blob<Dtype>*>* top);
+  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
+      vector<Blob<Dtype>*>* top);
+  virtual Dtype Backward_cpu(const vector<Blob<Dtype>*>& top,
+      const bool propagate_down, vector<Blob<Dtype>*>* bottom);
+  virtual Dtype Backward_gpu(const vector<Blob<Dtype>*>& top,
+      const bool propagate_down, vector<Blob<Dtype>*>* bottom);
+  int count_;
+};
+
+
 template <typename Dtype>
 class FlattenLayer : public Layer<Dtype> {
  public:

diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp
@@ -55,6 +55,8 @@ Layer<Dtype>* GetLayer(const LayerParameter& param) {
     return new SoftmaxLayer<Dtype>(param);
   } else if (type == "softmax_loss") {
     return new SoftmaxWithLossLayer<Dtype>(param);
+  } else if (type == "split") {
+    return new SplitLayer<Dtype>(param);
   } else if (type == "multinomial_logistic_loss") {
     return new MultinomialLogisticLossLayer<Dtype>(param);
   } else {