Caffeのカスタムレイヤーの練習

mnist_train1.prototxt

name: "mnist_train1"
layer {
  name: "mnist"
  type: "Data"
  top: "data"
  top: "label"
  include {
    phase: TRAIN
  }
  transform_param {
    scale: 0.00390625
  }
  data_param {
    source: "mnist_train_lmdb"
    batch_size: 100
    backend: LMDB
  }
}

layer {
  name: "mnist"
  type: "Data"
  top: "data"
  top: "label"
  include {
    phase: TEST
  }
  transform_param {
    scale: 0.00390625
  }
  data_param {
    source: "mnist_test_lmdb"
    batch_size: 100
    backend: LMDB
  }
}

layer {
  name: "conv1"
  type: "Python"
  bottom: "data"
  top: "conv1"
  python_param {
    module: 'my_conv_layer'
    layer: 'MyConvLayer'
    param_str: '{ "num_output": 20, "kernel_size": 5, "gaussian_std": 0.01}'
  }
}

layer {
  name: "relu1"
  type: "ReLU"
  bottom: "conv1"
  top: "relu1"
}

layer {
  name: "pool1"
  type: "Pooling"
  bottom: "relu1"
  top: "pool1"
  pooling_param {
    pool: MAX
    kernel_size: 2
    stride: 2
  }
}

layer {
  name: "conv2"
  type: "Python"
  bottom: "pool1"
  top: "conv2"
  python_param {
    module: 'my_conv_layer'
    layer: 'MyConvLayer'
    param_str: '{ "num_output": 50, "kernel_size": 5, "gaussian_std": 0.01}'
  }
}

layer {
  name: "relu2"
  type: "ReLU"
  bottom: "conv2"
  top: "relu2"
}

layer {
  name: "pool2"
  type: "Pooling"
  bottom: "relu2"
  top: "pool2"
  pooling_param {
    pool: MAX
    kernel_size: 2
    stride: 2
  }
}

layer {
  type: "Python"
  name: "fc3"
  top: "fc3"
  bottom: "pool2"
  python_param {
    module: 'my_inner_product_layer'
    layer: 'MyInnerProductLayer'
    param_str: '{"num_output": 500, "gaussian_std": 0.01}'
  }
}

layer {
  type: "ReLU"
  name: "relu3"
  top: "relu3"
  bottom: "fc3"
}

layer {
  type: "Python"
  name: "fc4"
  top: "fc4"
  bottom: "relu3"
  python_param {
    module: 'my_inner_product_layer'
    layer: 'MyInnerProductLayer'
    param_str: '{"num_output": 10, "gaussian_std": 0.01}'
  }
}

layer {
  name: "loss"
  type: "SoftmaxWithLoss"
  bottom: "fc4"
  bottom: "label"
  include: { phase: TRAIN }
}

layer {
  name: "accuracy"
  type: "Accuracy"
  bottom: "fc4"
  bottom: "label"
  top: "accuracy"
  include: { phase: TEST }
}

mnist_train1_solver.prototxt

net: "mnist_train1.prototxt"
test_iter: 100
test_interval: 500
base_lr: 0.5
lr_policy: "fixed"
display: 100
max_iter: 5000
snapshot_prefix: "mnist_train1"
solver_mode: CPU

my_conv_layer.py

import caffe
import numpy as np
import json

class MyConvLayer(caffe.Layer):
    def setup(self, bottom, top):
        assert len(bottom[0].shape) == 4
        j = json.loads(self.param_str)
        self._num_output = j['num_output']
        self._kernel_size = j['kernel_size']
        assert  self._kernel_size%2==1
        self._pad = (self._kernel_size-1)/2
        gaussian_std = j['gaussian_std']
        num_input = bottom[0].channels
        self.blobs.add_blob(self._num_output, num_input, self._kernel_size*self._kernel_size)
        self.blobs[0].data[...] = np.random.normal(0.0, gaussian_std, self.blobs[0].shape)

    def reshape(self, bottom, top):
        top[0].reshape(bottom[0].num, self._num_output, bottom[0].height, bottom[0].width)

    def _create_buffer(self, image):
        assert image.ndim == 3
        buf = np.zeros((self._kernel_size*self._kernel_size*image.shape[0], image.shape[1]*image.shape[2]), dtype=np.float32)
        col = 0
        for center_h in xrange(image.shape[1]):
            for center_w in xrange(image.shape[2]):
                row = 0
                for ch in xrange(image.shape[0]):
                    for h in xrange(center_h - self._pad, center_h + self._pad):
                        for w in xrange(center_w - self._pad, center_w + self._pad):
                            if 0 <= h < image.shape[1] and 0 <= w < image.shape[2]:
                                buf[row, col] = image[ch, h, w]
                            else:
                                buf[row, col] = 0.0
                            row += 1
                col += 1
        return buf

    def forward(self, bottom, top):
        weight_data = self.blobs[0].data.reshape(self._num_output,-1)
        top_data_flat = [weight_data.dot(self._create_buffer(image)) for image in bottom[0].data]
        top[0].data[...] = np.reshape(top_data_flat, top[0].shape)

    def backward(self, top, propagate_down, bottom):
        assert len(bottom[0].shape) == 4
        assert len(top[0].shape) == 4

        # update gradients
        self_diff_flat = [t.reshape(t.shape[0],-1).dot(self._create_buffer(b).T) for b,t in zip(bottom[0].data, top[0].diff)]
        self.blobs[0].diff[...] = np.sum(self_diff_flat, axis=0).reshape(self.blobs[0].diff.shape)

        # backward
        if propagate_down[0]:
            num_input = bottom[0].channels
            weight_data = self.blobs[0].data.transpose((1,0,2)).reshape(num_input,-1)
            bottom_diff_flat = [weight_data.dot(self._create_buffer(t)) for t in top[0].diff]
            bottom[0].diff[...] = np.reshape(bottom_diff_flat,bottom[0].shape)

my_inner_product_layer.py

import caffe
import numpy as np
import json

class MyInnerProductLayer(caffe.Layer):
    def setup(self, bottom, top):
        j = json.loads(self.param_str)
        self._num_output = j['num_output']
        gaussian_std = j['gaussian_std']
        self.blobs.add_blob(self._num_output, bottom[0].data[0].size)
        self.blobs[0].data[...] = np.random.normal(0.0, gaussian_std, self.blobs[0].shape)

    def reshape(self, bottom, top):
        top[0].reshape(bottom[0].num, self._num_output)

    def forward(self, bottom, top):
        b_data_flat = bottom[0].data.reshape((bottom[0].num,-1))
        top[0].data[...] = b_data_flat.dot(self.blobs[0].data.T)

    def backward(self, top, propagate_down, bottom):
        # update gradients
        b_data_flat = bottom[0].data.reshape((bottom[0].num,-1))
        self.blobs[0].diff[...] = top[0].diff.T.dot(b_data_flat)

        # backward
        if propagate_down[0]:
            t_diff_flat = top[0].diff.dot(self.blobs[0].data)
            bottom[0].diff[...] = t_diff_flat.reshape(bottom[0].shape)