Using torch.gather properly

MichaelMMeskhi: <p>I have a tensor of images of size <code>(3600, 32, 32, 3)</code> and I have a multi hot tensor [0, 1, 1, 0, …] of size <code>(3600, 1)</code>. I am looking to basically selecting images that correspond to a 1 in the multi hot tensor. I am trying to understand how to use torch.gather:</p>
<p><code>tensorA.gather(0, tensorB)</code></p>
<p>Gives me issues with dims and I can’t properly understand how to reshape them.</p>

ptrblck: <p>I’m not sure if I understand the use case correctly, but wouldn’t indexing the tensor with a mask work?</p>
<pre><code class="lang-python">x = torch.randn(3600, 32, 32, 3)
idx = torch.randint(0, 2, (3600,))
print(idx)
mask = idx.bool()

out = x[mask]
print(out.shape)
# torch.Size([1765, 32, 32, 3])
print(idx.sum())
# tensor(1765)
</code></pre>

MichaelMMeskhi: <p>The multi hot vector would be learned by the network. I believe doing <code>torch.gather</code> would allow for gradient updates. Correct me if I am wrong.</p>

ptrblck: <p>Yes, <code>gather</code> would not detach the computation graph, but I don’t understand how the model is learning these indices, as it seems that the multi-hot tensor would already be detached or how is your model returning integer indices without using e.g. <code>argmax</code>?</p>

MichaelMMeskhi: <p>Oh I see. Sorry for the confusing explanation. Let me rephrase it.</p>
<p>I have a Gumbel-Softmax network that outputs a vector <code>z</code> basically logits. I want to take non-zero indices in this <code>z</code> vector and select the same index location images from an image tensor that I have of 3600 images of 32x32x3 shape, tensor A <code>(3600, 32, 32, 3)</code>. But since my <code>z</code> is of shape <code>(1, 3600)</code>, I have been unable to do <code>.gather</code>. The output of gather should be something of the shape <code>(x, 32, 32, 3)</code> where <code>x</code> is the number of non-zero elements from <code>z</code>.</p>
<p>Later on, I want to optimize the Gumbel network to learn an optimal <code>z</code>.</p>

ptrblck: <p>Thanks for the follow up as this would match my understanding.<br>
However, what it unclear is how you are getting the integer indices (zeros and ones) from the logit output (<code>z</code>). Usually you would apply e.g. <code>torch.argmax</code> on the logits to create the indices and could then use it to index <code>A</code>. However, <code>argmax</code> would already detach the tensor so it wouldn’t matter if you are using my code snippet with a <code>mask</code>, <code>gather</code>, or any other operation.<br>
Assuming <code>z</code> contains logits (values between <code>[-Inf, +Inf]</code>) you won’t be able to use it as an index tensor.</p>

MichaelMMeskhi: <p>My understanding was I would do <code>torch.nonzero(z)</code> and then use that to do <code>gather</code>. To get one hot encoded I can use the <code>hard=True</code> flag <a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.gumbel_softmax.html" rel="noopener nofollow ugc">here</a>.</p>

ptrblck: <p>You are right about the <code>hard=True</code> flag in <code>gumbel_softmax</code>, but transforming the outputs via <code>nonzero()</code>, <code>out.long()</code>, or <code>out.bool()</code> would still detach it since the result would be an integer/bool type tensor, which is why I’m unsure how the selection could work (in a differentiable way).</p>
<pre><code class="lang-python">logits = torch.randn(20, requires_grad=True)
# Sample hard categorical using "Straight-through" trick:
out = F.gumbel_softmax(logits, tau=1, hard=True)
print(out.grad_fn)
# &lt;AddBackward0 object at 0x7f787c7d51f0&gt;

res = torch.nonzero(out)
print(res.grad_fn)
# None

print(res.long().grad_fn)
# None

print(res.bool().grad_fn)
# None
</code></pre>
<p>Could multiplying the output with <code>A</code> work somehow?</p>

MichaelMMeskhi: <p>I apologize I must have forgotten to mention. Using the <code>hard=True</code> flag, I am okay with using that output to as my indices to sample my images. I won’t need to do <code>argmax, nonzero, etc...</code></p>

ptrblck: <p>Unfortunately, you would need to use integer types:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(20)
x[out]
# IndexError: tensors used as indices must be long, byte or bool tensors
</code></pre>
<p>The same applies for gather:</p>
<blockquote>
<p><strong>index</strong> (<em>LongTensor</em>) – the indices of elements to gather</p>
</blockquote>

MichaelMMeskhi: <p>And type casting it as such would detach the gradients I assume?</p>

ptrblck: <p>No, floating point dtypes are fine, only integer types would detach it.<br>
I’m searching for some similar approaches in older posts to see if and how others have solved it.</p>

MichaelMMeskhi: <p>I see. Would a <code>torch.mv(imagesTensor, out)</code> work? Or something along those lines. Assuming <code>imagesTensor.shape = (3600, 32, 32, 3)</code> and <code>out.shape = (3600, 1)</code>.</p>

ptrblck: <p>Yeah, exactly. I was thinking if something like this could work for you:</p>
<pre><code class="lang-python">N = 10
x = torch.randn(N, 32, 32, 3)
out = torch.randint(0, 2, (N,)).float().requires_grad_(True)

res = x * out[:, None, None, None]

idx = res.nonzero()[:, 0].unique()
print(len(idx), out.sum())
# 5 tensor(5., grad_fn=&lt;SumBackward0&gt;)

res = res[idx]
print(res.shape)
# torch.Size([5, 32, 32, 3])

# perform your operation here
res = res.mean()

# backward
res.backward()

print(out)
# tensor([0., 0., 1., 1., 1., 1., 0., 0., 1., 0.], requires_grad=True)
print(out.grad)
# tensor([ 0.0000,  0.0000, -0.0044,  0.0005,  0.0032, -0.0044,  0.0000,  0.0000, 0.0046,  0.0000])
</code></pre>

MichaelMMeskhi: <p>Thank you so much! This is exactly what I was looking for. I appreciate you helping out.</p>