Sampling with index in 4-D tensor

1

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

Thanks

2
  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.
3

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

4

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