fast R-CNNとかなんとか。

バイトの関係でChainer用のRoIMaxPoolingを書いた。いろいろと雑。バグあるかも。

from chainer import function, cuda
from chainer.functions.array import concat

class RoIMaxPooling2D(function.Function):
    
    def __init__(self, pooled_shape, rects):
        self.kh = []
        self.kw = []
        self.rects = rects
        self.pooled_shape = pooled_shape
        for r in rects:
            self.kh.append(float(r[2]-r[0])/pooled_shape[0])
            self.kw.append(float(r[3]-r[1])/pooled_shape[1])
    
    def forward(self, x):
        xp = cuda.get_array_module(x)
        self.dim = x[0].shape[1]
        self.ys = []
        for r in range(len(self.rects)):
            y = xp.zeros((1, self.dim, self.pooled_shape[0], self.pooled_shape[1]))
            for i in range(self.pooled_shape[0]):
                for j in range(self.pooled_shape[1]):
                    y[ :, :, i, j] = xp.max(x[0][ :, :, xp.floor(i*self.kh[r]):xp.ceil((i+1)*self.kh[r]), xp.floor(j*self.kw[r]):xp.ceil((j+1)*self.kw[r])], axis=(2,3))
            self.ys.append(y.astype('f'))
        return concat.Concat(axis=0).forward(self.ys)
    
    def backward(self, x, gy):
        xp = cuda.get_array_module(x)
        self.gx = xp.zeros_like(x[0])
        for r in range(len(self.rects)):
            for i in range(self.pooled_shape[0]):
                for j in range(self.pooled_shape[1]):
                    self.gx[ :, :, xp.floor(i*self.kh[r]):xp.ceil((i+1)*self.kh[r]), xp.floor(j*self.kw[r]):xp.ceil((j+1)*self.kw[r])] += gy[0][r:r+1, :, i:i+1, j:j+1]
        return self.gx,

def roi_max_pooling_2d(x, pooled_shape, rects):
    return RoIMaxPooling2D(pooled_shape, rects)(x)