Caffe添加自己的Layer--Python

注：参考了实验室大佬的博客！
Caffe是部署性能比较好的深度学习开发框架，但是往往有时候现有的Layer不满足实际开发的要求，需要我们自定义一些Layer。Caffe添加Layer的方式主要有C++和Python两种。但需要注意的是，Caffe自定义的Python Layer不支持多GPU，只支持1个GPU训练。
一、编译Caffe源代码

1. 在编译Caffe时需要打开对Python Layer的支持：

gedit Makefile.config
# 然后将下面这行注释取消，保存
WITH_PYTHON_LAYER := 1
make all -j6 && make pycaffe

2. 添加环境变量

export LD_LIBRARY_PATH=`pwd`/.build_release/lib:$LD_LIBRARY_PATH
export PYTHONPATH=`pwd`/python:$PYTHONPATH

二、添加Python Layer
在python文件夹下面新建一个my_layer.py，主要包含了参数读入，Reshape，Forward和Backward等函数功能，内容如下：

#coding = utf-8
import numpy as np 
import caffe
class Mylayer(caffe.Layer):
	# 参数初始化
    def setup(self, bottom, top):
        if len(bottom) != 2:
            raise Exception("Need two inputs.")
        #参数解析
        params = eval(self.param_str)
        print(params['scale'])
	
	# data reshape
    def reshape(self, bottom, top):
        # bottom[0]: N x C x H x W
        # bottom[1]: 1 x 1 x 1 x M
        # top[0]: M x 1 x H x W
        self.num, self.channels, self.height, self.width = bottom[0].data.shape
        self.out_batch = bottom[1].data.shape[3]
        top[0].reshape(self.out_batch, 1, self.height, self.width)

    def forward(self, bottom, top):
        top[0].data[...] = 0
        label = bottom[1].data[0, 0, 0, :].copy()
        label = label.astype(np.int)

        for idx, lab in enumerate(label):
            top[0].data[idx, 0, :, :] = bottom[0].data[idx, lab, :, :]

    def backward(self, top, propagate_down, bottom):
        if propagate_down[1]:
            raise Exception("Can not backpropagate label.")
        if propagate_down[0]:
            bottom[0].diff[...] = 0
            label = bottom[1].data[0, 0, 0, :].copy()
            label = label.astype(np.int)

            for idx, lab in enumerate(label):
                bottom[0].diff[idx, lab] = top[0].diff[idx, 0, :, :]

三、测试Python Layer

1. 新建deploy.prototxt文件

name: "Mylayer"
input: "data"
input_dim: 10
input_dim: 128
input_dim: 20
input_dim: 20

input: "label"
input_dim: 1
input_dim: 1
input_dim: 1
input_dim: 32

layer {
    name: "mylayer"
    type: "Python"
    bottom: "data"
    bottom: "label"
    top: "translayer"
    python_param {
        module: "my_layer"
        layer: "Mylayer"
        param_str: "{\'scale\': 0.5}"
    }
}

2.使用caffe的一个工具进行测试，没有GPU的话去掉参数-gpu

./build/tools/caffe.bin time -model ./deploy.prototxt -gpu 0

四、在Python脚本中调用

# ....
net.data, net.label= .... # 数据输入或者前面层的结果
# ....
py_params = dict(scale=0.5)
layer_name = "mylayer"
net[layer_name] = L.Python(net.data, net.label, module="my_layer", layer="Mylayer", param_str=str(py_params))
# ....

注意：如果出现下面的错误:

layer_factory.hpp:77] Check failed: registry.count(type) == 1 (0 vs. 1) Unknown layer type: Python

则是编译时WITH_PYTHON_LAYER设置有问题，需要make clean，然后重新执行第一步