Caffe自定义网络层

参考博客

1 说明

本文描述了在Caffe源码中添加网络层的两种情况：不带配置参数的和带配置参数的。示例的网络层中没有做数据处理，传入的数据块直接传出去，也没有梯度计算。

2 不带参数的网络层

我们将要在 .prototxt 中添加的网络层为：

layer {
  name:"NewLayer"
  type:"New" # 网络类型
  bottom: "data"
  top:"data"
}

该网络层的类型名称为New，没有其他配置参数，只需要在对应位置添加 .cpp 和 .hpp 文件，添加步骤如下：

• 在 ./include/caffe/layers/ 文件夹下添加 new_layer.hpp：
  :::details

  #ifndef CAFFE_NEW_LAYER_HPP_
  #define CAFFE_NEW_LAYER_HPP_
   
  #include <vector>
   
  #include "caffe/blob.hpp"
  #include "caffe/layer.hpp"
  #include "caffe/proto/caffe.pb.h"
   
  namespace caffe {
  	template <typename Dtype>
  	class NewLayer : public Layer<Dtype> {
  	public:
  		explicit NewLayer(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){};
   
  	protected:
  		virtual void Forward_cpu(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 Forward_gpu(const vector<Blob<Dtype>*>& bottom,
  			const vector<Blob<Dtype>*>& top){};
  		virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
  			const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom){};
  	};
   
  } 
  #endif 
  

  :::

• 在 ./src/caffe/layers/ 文件夹下添加 new_layer.cpp：
  :::details

  #include <vector>
  #include "caffe/layers/new_layer.hpp"
   
  namespace caffe {
   
  template <typename Dtype>
  void NewLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
  	const vector<Blob<Dtype>*>& top) {
  	const Dtype* bottom_data = bottom[0]->cpu_data();
  	Dtype* top_data = top[0]->mutable_cpu_data();
  	const int count = bottom[0]->count();
  	for (int i = 0; i < count; ++i) {
  			top_data[i] = bottom_data[i];
  	}
  }
   
  template <typename Dtype>
  void NewLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
  	const vector<bool>& propagate_down,
  	const vector<Blob<Dtype>*>& bottom) {
  	if (propagate_down[0]) {
  		const Dtype* bottom_data = bottom[0]->cpu_data();
  		const Dtype* top_diff = top[0]->cpu_diff();
  		Dtype* bottom_diff = bottom[0]->mutable_cpu_diff();
  		const int count = bottom[0]->count();
  		for (int i = 0; i < count; ++i) {
  			bottom_diff[i] = top_diff[i];
  		}
  	}
  }
   
  #ifdef CPU_ONLY
  	STUB_GPU(NewLayer);
  #endif
   
  INSTANTIATE_CLASS(NewLayer);  //类名，注：这个类名与prototxt文件中的层名不需一致
  REGISTER_LAYER_CLASS(New); // 对应层的类型
   
  }  // namespace caffe
  

  :::

• 在 Caffe 的根目录下执行 make 即可编译。编译成功后即可在 .prototxt 中添加该网络层。

3 带参数的网络层

我们将要在 .prototxt 中添加的网络层为：

layer {
  name:"NewLayer"
  type:"New"
  bottom: "data"
  top:"data"
    new_param {
    coeff1: 1.0
	coeff2: 2.0
  }
}

相比于不带参数的网络层，带参数的网络层还需要 caffe.proto 文件中添加参数的说明。网络层添加步骤如下：

• 打开 ./src/caffe/proto/caffe.proto 文件

  • 在 message LayerParameter {} 中添加

      optional NewParameter new_param = 151;
    

  • 在其他空白出添加：

    message NewParameter{
        optional float coeff1 = 1 [default = 1];
        optional float coeff2 = 2 [default = 2];
    }
    

• 在 .cpp 中获取参数
  :::details

  #include <vector>
  #include "caffe/layers/new_layer.hpp"
   
  namespace caffe {
   
  template <typename Dtype>
  void NewLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
  	const vector<Blob<Dtype>*>& top) {
  	const Dtype* bottom_data = bottom[0]->cpu_data();
  	Dtype* top_data = top[0]->mutable_cpu_data();
  	const int count = bottom[0]->count();
  	
  	float coeff = this->layer_param_.new_param().coeff1(); // 获取参数
  	LOG(INFO) << "NewLayer, Forward_cpu:" << coeff;
   
  	for (int i = 0; i < count; ++i) {
  			top_data[i] = bottom_data[i];
  	}
  }
   
  template <typename Dtype>
  void NewLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
  	const vector<bool>& propagate_down,
  	const vector<Blob<Dtype>*>& bottom) {
   
  	float coeff = this->layer_param_.new_param().coeff2();	// 获取参数
  	LOG(INFO) << "NewLayer, Backward_cpu:" << coeff;
   
  	if (propagate_down[0]) {
  		const Dtype* bottom_data = bottom[0]->cpu_data();
  		const Dtype* top_diff = top[0]->cpu_diff();
  		Dtype* bottom_diff = bottom[0]->mutable_cpu_diff();
  		const int count = bottom[0]->count();
  		for (int i = 0; i < count; ++i) {
  			bottom_diff[i] = top_diff[i];
  		}
  	}
  }
   
  #ifdef CPU_ONLY
  	STUB_GPU(NewLayer);
  #endif
   
  INSTANTIATE_CLASS(NewLayer);  //类名，对应prototxt文件中的层名
  REGISTER_LAYER_CLASS(New); // 对应层的类型
   
  }  // namespace caffe
  

  :::

• 由于修改了 caffe.proto 文件，在执行 make 时会将 Caffe 重新编译一遍。

注意：添加的这些文件名、变量名要保持一致。