Maintaining gradient computation graph through input subset

username: I_M
cooked: <p>Hello,</p>
<p>I’m not sure how to best formulate this issue - I have a model that does something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, xyz):  # xyz is (B, 3, N)
    # a bunch of stuff that extract features
    # ....
    x = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))  # x: (B, 128, N) N= number of xyz points
</code></pre>
<p>Usually this would end in a semantic segmentation, but I’m not interested in that - I’m actually interesting in picking a <strong>unique</strong> subset of <strong>xyz</strong> with exactly a predetermined <code>num_points</code> - so the output would be shaped like <code>(B, 3, num_points)</code>.</p>
<p>I’m interested in using the learned extracted features, but I’m not sure how I can pick the subset of the points without disconnecting the gradient. One example of something I tried was:</p>
<pre data-code-wrap="python"><code class="lang-python"># self.weights = nn.Conv1d(128, 1, 1) 
scores = self.weights(x)  # # Shape: [B, 1, N]
scores = torch.softmax(weights, dim=-1) 

_, top_indices = torch.topk(weights.squeeze(1), self.num_points, dim=-1)  #  num points is the size of the subset
</code></pre>
<p>Using the indices I can then pick the subset, but this obviously disconnects the computational graph… Is there a way to sort of “force” the learned features to translate into a unique, fix-sized subset of the input features?</p>
<p>I hope I make sense. At the end, I expect to pass <code>model(xyz)</code> and get a unique hard subset of <code>xyz</code> shaped like <code>(B, 3, num_points)</code>.</p>
<p>Thank you for any tips <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="I_M" data-post="1" data-topic="208584">
<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/i_m/48/35728_2.png" class="avatar"> I_M:</div>
<blockquote>
<p>Using the indices I can then pick the subset, but this obviously disconnects the computational graph…</p>
</blockquote>
</aside>
<p>This isn’t the case if you select from a differentiable output:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(1, 1)
lin = nn.Linear(1, 10)

out = lin(x)

out_selected, indices = torch.topk(out, k=3, dim=1)
print(out_selected.shape)
# torch.Size([1, 3])
print(out_selected.grad_fn)
# &lt;TopkBackward0 object at 0x7f4ce3ec5210&gt;
print(indices)
# tensor([[1, 7, 0]])

out_selected.mean().backward()
print(lin.weight.grad)
# tensor([[-0.0697],
#         [-0.0697],
#         [ 0.0000],
#         [ 0.0000],
#         [ 0.0000],
#         [ 0.0000],
#         [ 0.0000],
#         [-0.0697],
#         [ 0.0000],
#         [ 0.0000]])
</code></pre>
<p>In your example code you are also calling <code>torch.topk(weights...)</code> while in your initial example you want to process the <code>xyz</code> input to the model, not the trainable <code>weight</code>, so unsure what exactly your use case is.</p>

username: I_M
cooked: <p>Hey ptrblck!<br>
Thanks for the reply.</p>
<p>I meant using the <strong>indices</strong>, rather than the <code>out_selected</code> as in your reply<br>
I’m not sure if my goal here is actually feasible, perhaps you would know a little better than me.</p>
<p>I’ll try explaining my expected input/output</p>
<p><strong>Input:</strong> 3D points, lets call these <code>xyz</code><br>
<strong>Output:</strong> Exactly <code>n</code> 3D points which are a <strong>hard &amp; unique</strong> subset of <code>xyz</code>, lets call these <code>out</code><br>
This is basically what I want the model to learn.</p>
<p>My model extracts features from the xyz points through various techniques. For clarity, lets call these <code>features</code> and they are shaped <code>(B, 128, N)</code>.</p>
<p>I want to use <code>features</code> in combination with the input <code>xyz</code> to produce <code>out</code>.</p>
<p>Edit:<br>
When I mentioned the “trick” that fails, with the weights layer - I attempted to do something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">scores = self.weights(features)  # # self.weights = nn.Conv1d(128, 1, 1) 
scores = scores.squeeze(1)

_, indices = torch.topk(F.softmax(scores, dim=1), 30, dim=-1, largest=True, sorted=True)
selected_points = torch.gather(xyz, 2, indices.unsqueeze(1).expand(-1, 3, -1))  # (B, 3, 30) - This is "out"
</code></pre>

username: ptrblck
cooked: <p>I’m still missing details of your use case since <code>selected_points</code> are still differentiable:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.functional as F

weights = nn.Conv1d(128, 1, 1) 
features = torch.randn(1, 128, 100)

scores = weights(features)  # # self.weights = nn.Conv1d(128, 1, 1) 
scores = scores.squeeze(1)

_, indices = torch.topk(F.softmax(scores, dim=1), 30, dim=-1, largest=True, sorted=True)

xyz = torch.randn(1, 3, 100, requires_grad=True)
selected_points = torch.gather(xyz, 2, indices.unsqueeze(1).expand(-1, 3, -1))  # (B, 3, 30) - This is "out"

print(selected_points.shape)
# torch.Size([1, 3, 30])
print(selected_points.grad_fn)
# &lt;GatherBackward0 object at 0x7ff13ddf9d80&gt;

selected_points.mean().backward()

print(xyz.grad)
# tensor([[[0.0000, 0.0111, 0.0111, 0.0000, 0.0000, 0.0000, 0.0111, 0.0111,
#           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0111, 0.0111, 0.0000,
#           ...
</code></pre>
<p>and I don’t know if you want to <em>learn</em> the <code>indices</code>, which is not directly possible since integer values would have a gradient of zero everywhere and Inf/NaN at the rounding points.</p>