在Caffe2中添加自定义的Python层

3
在Caffe中添加Python层相当简单(创建一个从caffe.layer继承的子类并添加四个基本方法,如此处此处所述)。然而,在Caffe2中添加自定义Python层对我来说并不那么简单。有人能解释一下在Caffe2中添加Python层的步骤吗?

一个Caffe2问题,涉及到一个可能有帮助的示例: https://github.com/caffe2/caffe2/issues/366 - rkellerm
1个回答

0

首先,您必须按照示例将新层实现为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.]


网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接