Select elements of Tensor based on index tensor along same dimensions

user_7484221 0
Viewed 668

I have the following two tensors

input shape: 16 32 32 3
index shape: 16 32 32 2
output shape: 16 32 32 3

The formula for the output would be:

output[b, h, w] = input[b, index[b, h, w, 0], index[b, h, w, 1]]

I tried to use torch.gather but I was not able to formulate the previous assignment.

Does anyone know how to do this in an efficient manner? Thanks!

For context: input contains a batch of 16 elemens where each one is a tensor of 32x32 that containts 3D points. index is a mapping from position to 3D point.

pytorch python
1 Answers

You can achieve this by unraveling the indices from index (on dimensions 1 and 2) in order to index input on a single dimension using torch.gather.

This requires to expand the shape of the indexer to fit the shape of input:

Here is an example with some dummy data:

>>> x = torch.rand(16, 32, 32, 3)
>>> index = torch.randint(0, 10, (16,32,32,2))

Some manipulation on index is required to unravel the values:

>>> unraveled = x.size(1)*index[..., 0] + index[..., 1]
>>> u = unraveled.flatten(1).unsqueeze(-1).expand(-1, -1, x.size(-1))

Now u, reshaped expanded version of index has a shape of (16, 1024, 3).

The indexed tensor also needs to be flattened:

>>> x.flatten(1, 2)
torch.Size([16, 1024, 3])

Finally, you can gather on dim=1 (keep in mind the result needs to be reshaped to the desired shape i.e. the input's shape):

>>> out = input.flatten(1,2).gather(1, u).reshape_as(x)
user_63313690
Related