Indexing Tensor

#1

I have two tensors and want to access one of them using the other one as index.
a.shape = [batch_size, num_channels, height, width]
b.shape = [batch_size, num_channels, 2]

I would now like to use the 2 as x and y coordinates of the feature maps of a, but I don’t know how to do that… Can anyone give me a hint?
The expected output would be of size [batch_size, num_channels]

Thanks in advance

#2

I guess, I could do it by expanding the indices and calling gather twice, as follows:

x = b[:,:,0]
x = x.view(4,16,1,1).expand(4,16,256,1)
y = b[:,:,1]
y = y.view(4,16,1,1)
result = a.gather(3, x).gather(2, y).squeeze()

but I think this looks a bit clumsy and I feel like there should be a better way to accomplish this.

#3

Hi,

I’m afraid there is not very elegant way to do this.
Another approach would be to convert your 2D indices into a 1D index into a dimension of size height * width.
Then you can use view of a and a single gather to get what you want.

#4

Thanks for your reply. Is there an elegant way to convert the indices to 1D?

#5

You can do the following:

import torch

batch_size = 2
num_channels = 3
height = 5
width = 5

a = torch.rand(batch_size, num_channels, height, width)
b = torch.rand(batch_size, num_channels, 2).mul(min(height, width)).long() # Make random valid indices

# Your solution
x = b[:,:,0]
x = x.view(batch_size,num_channels,1,1).expand(batch_size,num_channels,height,1)
y = b[:,:,1]
y = y.view(batch_size,num_channels,1,1)
result = a.gather(3, x).gather(2, y).squeeze()

# 1D solution
lin_indices = b.select(-1, 1) * width + b.select(-1, 0)
assert a.is_contiguous()
lin_a = a.view(batch_size, num_channels, height*width)
result2 = lin_a.gather(-1, lin_indices.unsqueeze(-1)).squeeze(-1)

print((result - result2).abs().max())