在caffe中添加python层非常简单(创建一个继承自caffe.layer并添加四种基本方法的子类,如here和here所述。但是,在caffe2中添加自定义python层对我来说并不是那么简单。请解释在caffe2中添加python层的过程?
首先,您必须将新层实现为Python类,如示例所示。在这种情况下,它只以相反的顺序输出输入张量:
class ReverseOrderOp(object):
def forward(self, inputs, outputs):
blob_out = outputs[0]
blob_out.reshape(inputs[0].shape)
blob_out.data[...] = inputs[0].data[::-1]
然后,您可以使用model.net.Python
将新图层添加到模型中:
model = ModelHelper(name="test")
l = np.asarray([0,1,2,3])
workspace.FeedBlob('l', l.astype(np.float32))
model.net.Python(ReverseOrderOp().forward)(
['l'], ['out'], name='ReverseOrder'
)
workspace.RunNetOnce(model.net)
print(workspace.FetchBlob('out'))
输出应该是[ 3. 2. 1. 0.]