How to use array as indices to take values from another array

Manuel_Alejandro_Dia: <p>Hi,</p>
<p>I am coding a Monocular Depth Perception Network that follows a bin + residual process. I have a set number of bins which represent a range of values between two numbers and are defined by their left and right edges, and a midpoints (these last ones are encoded in logarithmic space).</p>
<p>They are defined as:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/e/2e3b94d074db89d0646b091fa3a081526ed2b0f0.png" data-download-href="https://discuss.pytorch.org/uploads/default/2e3b94d074db89d0646b091fa3a081526ed2b0f0" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e3b94d074db89d0646b091fa3a081526ed2b0f0_2_543x500.png" alt="image" data-base62-sha1="6AZCbTWX56BSRUb0rOUbYJ9jZh6" width="543" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e3b94d074db89d0646b091fa3a081526ed2b0f0_2_543x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e3b94d074db89d0646b091fa3a081526ed2b0f0_2_814x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/e/2e3b94d074db89d0646b091fa3a081526ed2b0f0.png 2x" data-dominant-color="E6E6E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">966×888 127 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/4/1462e2e294bc78136d462e7470e9d0107a6a4546.png" data-download-href="https://discuss.pytorch.org/uploads/default/1462e2e294bc78136d462e7470e9d0107a6a4546" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/4/1462e2e294bc78136d462e7470e9d0107a6a4546.png" alt="image" data-base62-sha1="2UlpyKiIalWUjLarn0DtUAIPPue" width="543" height="500" data-dominant-color="E9E9E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">983×305 35.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">import torch
import math

min_depth=1.
max_depth=126.
num_bins=48

alpha = min_depth
beta = max_depth
bins_edges = [math.exp( math.log(alpha) + ( (math.log(beta/alpha)*i) / num_bins ) ) for i in range(num_bins+1)]
midpoints = [ ( math.log(bins_edges[i+1]) + math.log(bins_edges[i]) ) / 2 for i in range(num_bins)]
bins_edges = torch.tensor(bins_edges)
midpoints_log = torch.tensor(midpoints)
</code></pre>
<p>My network predicts two feature maps, one that I call <code>bin_logits</code> which tells me the probability that a pixel has a depth value within a bin (classification problem), and another called <code>residuals</code> which helps me give an extra value to the depth so it is continuous together with the midpoint defined for the bin using the following formula:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/b/6bf52018700efd9b1c3159cd076ff8a30f5b3436.png" data-download-href="https://discuss.pytorch.org/uploads/default/6bf52018700efd9b1c3159cd076ff8a30f5b3436" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/b/6bf52018700efd9b1c3159cd076ff8a30f5b3436.png" alt="image" data-base62-sha1="fp2gCiqroYDkUx1hJlNPmU1AgBg" width="543" height="500" data-dominant-color="E4E4E4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">986×211 29.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Which means that depending on the bins predicted by <code>bin_logits</code>, I need to take the respective value from the from the <code>K</code> layer in the <code>residuals</code> feature map for each pixel. I need to also take the respective midpoint, left and right edges in order to calculate the depth.</p>
<p>I have been able to use the indexes given by <code>argmax</code> on the <code>bin_logits</code> features to take the midpoints, and also the left and right edges like:</p>
<pre><code class="lang-auto">#dummy output for a batch of 8 images of size 100x200
bin_logits = torch.rand((8,num_bins,100,200))
residuals = torch.rand((8,num_bins,100,200))

bin_logits = bin_logits.argmax(1)
midpoints = midpoints_log[bin_logits.flatten()].view_as(bin_logits)
left_edges = bins_edges[bin_logits.flatten()].view_as(bin_logits)
right_edges = bins_edges[(bin_logits+1).flatten()].view_as(bin_logits) 
</code></pre>
<p>But I do not know how to do it for the <code>residuals</code> feature map. I tried using <code>bin_logits.max(dim=1, keep_dim=True)</code>, but it gives me indices in the range of the number of bins (0-47 in this case).</p>
<p>I need for the final tensor to be of shape <code>(8,100,200)</code>. So I can apply the equation <code>(4)</code> seen above.</p>
<p><strong>How can I use the</strong> <code>argmax</code> <strong>from</strong> <code>bin_logits</code> <strong>feature maps to select the corresponding values in the</strong> <code>residuals</code> <strong>tensor?</strong> Like I do for the edges and midpoint.</p>
<p>Thanks for any help!</p>

Manuel_Alejandro_Dia: <p>I found the answer using <code>torch.gather</code>:</p>
<pre><code class="lang-auto">residuals = torch.gather(residuals, 1, bin_logits.unsqueeze(1)).squeeze(1)
</code></pre>
<p>What I needed to do was to make the number of dimensions for the index array and the residuals equal. Then collapse the expanded dimension, et voilá!</p>