caffe自定義層

1. 實作hpp和cpp

net在拿到layer之後會調用每層layer的Setup函數,每層Setup中會調用：

void SetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
    CheckBlobCounts(bottom, top);
    LayerSetUp(bottom, top);
    Reshape(bottom, top);
    SetLossWeights(top);
}
           

其中cpp中要重載四個函數

• LayerSetup
• Reshape
• Forward
• Backward

upsample.h

#ifndef CAFFE_UPSAMPLE_LAYER_HPP_
#define CAFFE_UPSAMPLE_LAYER_HPP_
 
#include <vector>
 
#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"
 
namespace caffe {
 
template <typename Dtype>
class UpsampleLayer : public Layer<Dtype> {
 public:
  explicit UpsampleLayer(const LayerParameter& param)
      : Layer<Dtype>(param) {}  //構造
  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
 
  virtual inline const char* type() const { return "Upsample"; } //名字
  virtual inline int MinBottomBlobs() const { return 1; }    //blobs數量限制
  virtual inline int MaxBottomBlobs() const { return 1; }
  virtual inline int ExactNumTopBlobs() const { return 1; }
 
  //實際要重載這四個函數，但是父類的gpu預設調用cpu傳回，如果不寫cuda則隻需要重載cpu即可
  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top);
  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
 
 private:
  int scale_;
};
 
 
 
} // namespace caffe
 
#endif // CAFFE_UPSAMPLE_LAYER_HPP_
           

upsample.cpp

#include <vector>
#include "caffe/layers/upsample_layer.hpp"
 
namespace caffe {
 
template <typename Dtype>
void UpsampleLayer<Dtype>::LayerSetUp(
  const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
  UpsampleParameter upsample_param = this->layer_param_.upsample_param(); //UpsampleParameter 類會由caffe.proto注冊後自動生成
  scale_ = upsample_param.scale(); 
}
 
template <typename Dtype>
void UpsampleLayer<Dtype>::Reshape(  //reshape的作用是一緻bottom為top開空間，在net定義的時候調用
  const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
  vector<int> out_shape;
  for (int i = 0; i < bottom[0]->num_axes(); i++) {
    out_shape.push_back(bottom[0]->shape(i));
  }
 
  out_shape[bottom[0]->num_axes() - 1] *= scale_;
  out_shape[bottom[0]->num_axes() - 2] *= scale_;
  top[0]->Reshape(out_shape);
}

template <typename Dtype>
void UpsampleLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) { 
 
  int N = top[0]->shape(0);
  int C = top[0]->shape(1);
  int H = top[0]->shape(2);
  int W = top[0]->shape(3);
 
  const Dtype *input = bottom[0]->cpu_data();
  Dtype *output = top[0]->mutable_cpu_data();
  for (int n = 0; n < N; n++) {
    for (int c = 0; c < C; c++) {
      for (int h = 0; h < H; h++) {
        for (int w = 0; w < W; w++) {
          int nw = w/scale_;
          int nh = h/scale_;
          int out_idx = (((n * C + c) * H) + h) * W + w;
          int in_idx = (((n * C + c) * (H / scale_)) + nh) * (W / scale_) + nw;
          output[out_idx] = input[in_idx];
        }
      }
    }
  }
}
 
template <typename Dtype>
void UpsampleLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  int N = bottom[0]->shape(0);
  int C = bottom[0]->shape(1);
  int H = bottom[0]->shape(2);
  int W = bottom[0]->shape(3);
  const Dtype *output_grad = top[0]->cpu_diff();
  Dtype* bottom_diff = bottom[0]->mutable_cpu_diff();
  caffe_set(bottom[0]->count(), Dtype(0), bottom_diff);
  for (int n = 0; n < N; n++) {
    for (int c = 0; c < C; c++) {
      for (int h = 0; h < H; h++) {
        for (int w = 0; w < W; w++) {
          for (int i = 0; i < scale_; i++) {
            for (int j = 0; j < scale_; j++) {
              int nw = w * scale_ + i;
              int nh = h * scale_ + j;
              int out_idx = (((n * C + c) * H) + h) * W + w;
              int in_idx = (((n * C + c) * (H * scale_))
                  + nh) * (W * scale_) + nw;
              bottom_diff[out_idx] += output_grad[in_idx];
            }
          }
        }
      }
    }
  }
}
 
#ifdef CPU_ONLY
STUB_GPU(UpsampleLayer);
#endif
 
INSTANTIATE_CLASS(UpsampleLayer); //執行個體化float和double的layer
REGISTER_LAYER_CLASS(Upsample);   //在layer_factory裡生成層的float和double creator
 
} // namespace caffe

           

1. 修改caffe.prorotxt

• 給予編号，optional xxxParameter xxx_param = 152
• 添加message xxxParameter{}

message AllPassParameter {  

 optional float key = 1 [default = 0];  

}
           

重編 即可