Ajouter une couche Python personnalisée dans Caffe2

Question

Ajouter une couche python dans Caffe était assez simple (en créant une classe enfant qui hérite de caffe.layer et en ajoutant quatre méthodes de base, comme décrit ici et ici). Cependant, ajouter une couche python personnalisée dans Caffe2 n'est pas aussi simple pour moi. Quelqu'un peut-il s'il vous plaît expliquer la procédure pour ajouter une couche python dans Caffe2?

Answer 1

Tout d'abord, vous devez implémenter votre nouvelle couche en tant que classe Python comme indiqué dans l'exemple. Dans ce cas, elle ne fait que renvoyer le tenseur d'entrée dans l'ordre inverse :

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]

Ensuite, vous pouvez ajouter votre nouvelle couche au modèle en utilisant 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'))

La sortie devrait être [ 3. 2. 1. 0.]