How to select indices from a 4D tensor?

1

I have a feature map of size NxCxWxH, where,

  • N is the batch size
  • C number of channels
  • W width
  • H height

I have indices of length N where each index in this list corresponds to a channel. I need to select the channels based on these indices. The size of the output should be Nx1xWxH or NxWxH. How can I do this?

2

Hi,

I think you want to expand your indices as Nx1xWxH by doing ind.view(N, 1, 1, 1).expand(N, 1, W, H).
Then you can use gather to get the values: res = value.gather(ind, 1).
Then if you want to remove the dimension of size 1, you can do res = res.squeeze(1).

1 Like
3

Thank you so much for your reply. It solved my problem.

4

Hi Alban (and smani)!

How about direct indexing?

import torch
torch.__version__

torch.manual_seed (2020)

N = 5
C = 4
W = 3

t = torch.randn (N, C, W)
t

i = torch.LongTensor ([3, 0, 3, 1, 2])
i

t[torch.arange (0, N).long(), i, :]

>>> import torch
>>> torch.__version__
'0.3.0b0+591e73e'
>>>
>>> torch.manual_seed (2020)
<torch._C.Generator object at 0x00000274E0EF6630>
>>>
>>> N = 5
>>> C = 4
>>> W = 3
>>>
>>> t = torch.randn (N, C, W)
>>> t

(0 ,.,.) =
  1.2372 -0.9604  1.5415
 -0.4079  0.8806  0.0529
  0.0751  0.4777 -0.6759
 -2.1489 -1.1463 -0.2720

(1 ,.,.) =
  1.0066 -0.0416 -1.2853
 -0.4948 -1.2964 -1.2502
 -0.7693  1.6856  0.3546
  0.7790  0.3257  0.4995

(2 ,.,.) =
  0.7705 -0.5920  0.5270
  0.0807  0.9863  2.2251
 -1.1789 -0.3879 -0.3383
 -1.3321  0.2644 -0.4644

(3 ,.,.) =
  1.5745 -0.7463  0.5680
  0.5286  1.9171 -0.2539
 -1.2216 -0.0740 -1.8290
 -0.4015 -0.2322  0.2829

(4 ,.,.) =
  1.4925  0.7024 -1.5135
  0.1471  0.5943 -2.1462
  0.2443  0.4215 -0.4506
 -3.1644 -0.3315 -0.2427
[torch.FloatTensor of size 5x4x3]

>>>
>>> i = torch.LongTensor ([3, 0, 3, 1, 2])
>>> i

 3
 0
 3
 1
 2
[torch.LongTensor of size 5]

>>>
>>> t[torch.arange (0, N).long(), i, :]

-2.1489 -1.1463 -0.2720
 1.0066 -0.0416 -1.2853
-1.3321  0.2644 -0.4644
 0.5286  1.9171 -0.2539
 0.2443  0.4215 -0.4506
[torch.FloatTensor of size 5x3]

Best.

K. Frank

5

That works as well.
Not sure which one will be fastest and use less memory…

My personal preference is usually to avoid fancy advanced indexing because I’m not used to that semantic :smiley: Also allocating the arange() Tensor can be expensive depending on the sizes involved. Note that expand does not allocate memory so my solution don’t do any extra allocation.