How to select indices from a 4D tensor?

smani: <p>I have a feature map of size NxCxWxH, where,</p>
<ul>
<li>N is the batch size</li>
<li>C number of channels</li>
<li>W width</li>
<li>H height</li>
</ul>
<p>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?</p>

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

smani: <p>Thank you so much for your reply. It solved my problem.</p>

KFrank: <p>Hi Alban (and smani)!</p>
<aside class="quote no-group" data-username="albanD" data-post="2" data-topic="79243" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>I think you want to expand your indices as Nx1xWxH by doing <code>ind.view(N, 1, 1, 1).expand(N, 1, W, H)</code>.<br>
Then you can use gather to get the values: <code>res = value.gather(ind, 1)</code>.</p>
</blockquote>
</aside>
<p>How about direct indexing?</p>
<pre data-code-wrap="python"><code class="lang-python">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, :]
</code></pre>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; torch.manual_seed (2020)
&lt;torch._C.Generator object at 0x00000274E0EF6630&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; N = 5
&gt;&gt;&gt; C = 4
&gt;&gt;&gt; W = 3
&gt;&gt;&gt;
&gt;&gt;&gt; t = torch.randn (N, C, W)
&gt;&gt;&gt; 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]

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

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

&gt;&gt;&gt;
&gt;&gt;&gt; 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]
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

albanD: <p>That works as well.<br>
Not sure which one will be fastest and use less memory…</p>
<p>My personal preference is usually to avoid fancy advanced indexing because I’m not used to that semantic <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"> Also allocating the <code>arange()</code> 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.</p>