Использование PythonLayer в caffe

[BeS 53]

Всем привет,

Решил попробовать пистаь к caffe кастомные слои не на C++, а при помощи имеющего PythonLayer из стандартного caffe.

Пишу простой сэмпл:

1) так выглядит sample.prototxt

Quote

name: "sample"
input: "data"
input_dim: 1
input_dim: 3
input_dim: 1280
input_dim: 854
layer {
  name: 'sample_layer'
  type: 'Python'
  top: 'sample_layer'
  bottom: 'data'
  python_param {
    module: 'sample_layer'
    layer: 'SampleLayer'
  }
}
layer {
   type: "Pooling"
   name: "pool"
   bottom: "sample_layer"
   top: "pool"
   pooling_param {
     pool: MAX
     kernel_size: 2
     stride: 2
   }
}
 

2) так выглядит собственно код Python-слоя (должен тупо пропихивать блоб с данными дальше без изменения)

Quote

import caffe
import numpy as np
import yaml

class SampleLayer(caffe.Layer):
    def setup(self, bottom, top):
        pass

    def reshape(self, bottom, top):
        pass

    def forward(self, bottom, top):
        top[0].reshape(*bottom[0].shape)
        top[0].data[...] = bottom[0].data

    def backward(self, top, propagate_down, bottom):
        pass
 

3) так выглядит код, который запускает сетку

Quote

import numpy as np
import caffe, cv2

caffe.set_mode_cpu()

net = caffe.Net('sample.prototxt',caffe.TEST)
im = cv2.imread("img.jpg")
height, width, channels = im.shape
net.blobs['data'].reshape(1, channels, width, height)
new_im = im.swapaxes(0, 2)
net.blobs['data'].data[0] = new_im
net.forward()

Какое поведение наблюдается при запуске сети:

1) Отключаю пулинг слой, сеть работает корректно. После forward pass я могу с выхода кастомного слоя забрать картинку в первозданном виде.

2) Отключаю python-слой а в качестве bootm для Python-слоя задаю просто data. Все работает корректно и на выходе получаю картинку вдвое меньшую, чем подавал на вход.

3) Запускаю сеть в которой сначала идет мой кастомный слой, а его выход идет на пулинг. Сеть не работает, ругаясь на невалидный блоб на входе пулинга:

Quote

F0424 13:13:18.812511  3544 pooling_layer.cpp:84] Check failed: 4 == bottom[0]->num_axes() (4 vs. 0) Input must have 4 axes, corresponding to (num, channels, height, width)
*** Check failure stack trace: ***
Aborted (core dumped)

Может кто-то сталкивался с таким и сумел пофиксить?