how to apply "Gather" operation like numpy in Caffe2?

I am new to Caffe2, and I want to compose an operation like this:

Numpy way

pytoch way

My question is, how to compose Caffe2 operators to make the same operators like above? I have tried some compositions but still I couldn't find the right one. If anyone knows the composition, please help, I will be really appreciate for it.

Solution

There is a Gather operator in Caffe2. The main problem with this operator is that you can't set the axis (it's always 0). So, if we run this code:

model = ModelHelper(name="test")

s = np.arange(20).reshape(4, 5)
y = np.asarray([0, 1, 2])

workspace.FeedBlob('s', s.astype(np.float32))
workspace.FeedBlob('y', y.astype(np.int32))

model.net.Gather(['s', 'y'], ['out'])

workspace.RunNetOnce(model.net)

out = workspace.FetchBlob('out')
print(out)

We will get:

[[  0.   1.   2.   3.   4.]
 [  5.   6.   7.   8.   9.]
 [ 10.  11.  12.  13.  14.]]

One solution could be to reshape s to a 1D array and transform y in the same way. First of all, we have to implement an operator to transform y. In this case, we will use a numpy function called ravel_multi_index:

class RavelMultiIndexOp(object):
    def forward(self, inputs, outputs):
        blob_out = outputs[0]

        index = np.ravel_multi_index(inputs[0].data, inputs[1].shape)

        blob_out.reshape(index.shape)
        blob_out.data[...] = index

Now, we can reimplement our original code:

model = ModelHelper(name="test")

s = np.arange(20).reshape(4, 5)
y = np.asarray([[0, 1, 2],[0, 1, 2]])

workspace.FeedBlob('s', s.astype(np.float32))
workspace.FeedBlob('y', y.astype(np.int32))

model.net.Python(RavelMultiIndexOp().forward)(
    ['y', 's'], ['y'], name='RavelMultiIndex'
)
model.net.Reshape('s', ['s_reshaped', 's_old'], shape=(-1, 1))

model.net.Gather(['s_reshaped', 'y'], ['out'])

workspace.RunNetOnce(model.net)

out = workspace.FetchBlob('out')
print(out)

Output:

[[  0.]
 [  6.]
 [ 12.]]

You may want to reshape it to (1, -1).