Sampling with index in 4-D tensor

I have a pair of value - index tensors.

value = torch.rand(2, 128, 10, 10) # (batch, feature_size, x, y)
indices = torch.randint_like(torch.zeros(2, 30, 2), min=0, max=9) # (batch, n_sample, xy)

The expected shape of output tensor is like below,

res: torch.Size([2, 128, 30]) # (batch, feature_size, n_sample)

This operation means sampling 30 points among the 10x10 = 100 plane.
I tried to use index_select, but it only works with index tensor of 1-D.
I have an option that create a bool tensor having size of (2, 10, 10) and use masked_select.
However, it is ineffective as the number of samples increase.

I need your help


  1. You’re selecting 30 samples per plane from 128 planes, all samples at the same position per batch. Why not just use slices across planes?
  2. With batch_num==2 you can iterate IMO, and also you can just flatten the 10x10 into 100 if 4 dims is hard to keep in mind.

I couldn’t get it ‘all samples at the same position per batch’ from the first comment.
But flattening a tensor to easily handle is helpful.

Thank you for your advices

I mean you sample 30 examples per plane. Each of examples is located at the same coords [xi,yi]. Your script is generating 30 different positions per batch, not per each plane. Thus, if this is not a mistake, you can slice through planes to retrieve 30 stacks of 1x128, since sample[i] is always at [xi,yi]

1 Like