Adding a custom python layer in caffe2

Question

Adding a python layer in caffe was fairly straightforward (creating a child class that inherits from caffe.layer and adding four basic methods, as described here and here. However, adding a custom python layer in caffe2 is not as straightforward to me. Can someone please explain the procedure for adding a python layer in caffe2?

Answer 1

First, you must implement your new layer as a Python class as shown in the example. In this case, it only outputs the input tensor in reverse order:

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]

Then, you can add your new layer to the model using 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'))

The output should be [ 3. 2. 1. 0.]