Creating an output tensor while preserving autograd graph

Joseph_M: <p>Hello!<br>
I’m trying to implement a MaxPooling2D layer from scratch without using any modules from <code>toch.nn</code>.</p>
<pre><code class="lang-auto">def forward(self, x, debug=False):
        # In shape
        batch, in_features, h_in, w_in = x.shape
        # Expected out dims
        # If the img has been padded the term + 2 * padding should not be considered
        h_out = int(((h_in - self.kernel) / self.stride)) + 1
        w_out = int(((w_in - self.kernel) / self.stride)) + 1
        # Preparing output
        out = torch.zeros(size=(batch, in_features, h_out * w_out))
        # Getting all the windows (basically what fold does)
        offset = 0
        for i in range(0, h_in, self.stride):
            for j in range(0, w_in, self.stride):
                if j + self.kernel &gt; w_in or i + self.kernel &gt; h_in:
                    continue
                window_elems = x[:, :, i:i + self.kernel, j:j + self.kernel].long()
                max_elems = torch.amax(window_elems, in_features)
                elems = torch.amax(max_elems, 2)
                out[:, :, offset] = elems
                offset += 1
        return out.reshape(batch, in_features, h_out, w_out)
</code></pre>
<p>The problem that I’m currently having is with the out tensor that not only is considered a leaf node but also it has no grad_fn (and that means I can not call the backward from the model). Is there any way that I can create the out vector while “linking it” with the input vector in the autograd graph?<br>
Thanks in advance!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/joseph_m">@Joseph_M</a>,</p>
<p>You could try adding <code>0*input</code>, and see if that works? The output does have a different shape than the input but you might be able to connect them via some trick like multiplying by 0 and adding it to the output!</p>

ptrblck: <aside class="quote no-group" data-username="Joseph_M" data-post="1" data-topic="179986">
<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/joseph_m/48/58460_2.png" class="avatar"> Joseph_M:</div>
<blockquote>
<p>I’m trying to implement a MaxPooling2D layer from scratch without using any modules from <code>toch.nn</code>.</p>
</blockquote>
</aside>
<p>It seems you are assigning the <code>amax</code> indices to the output tensors, which will be detached from the computation graph since <code>amax</code> is not differentiable, instead of the <code>max</code> values, which differs from the <code>nn.MaxPool2d</code> implementation:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(1, 1, 2, 2, requires_grad=True)
pool = nn.MaxPool2d(2, return_indices=True)
out, idx = pool(x)
print(out)
# tensor([[[[-0.0955]]]], grad_fn=&lt;MaxPool2DWithIndicesBackward0&gt;)
print(idx)
# tensor([[[[2]]]])
</code></pre>
<p>As you can see, the actual <code>out</code> tensor has a valid <code>grad_fn</code> while the corresponding <code>idx</code> does not.</p>

Joseph_M: <p>Thanks for the reply and for the insight! I have just a quick doubt, you said amax is not differentiable but if I try the following code:</p>
<pre><code class="lang-auto">x = torch.randint(0,255,size=(1,3,20,20),dtype=torch.float32,requires_grad=True)
y = torch.amax(x)
print(y.requires_grad,y.grad_fn)
# True &lt;AmaxBackward0 object at 0x00000184988F3190&gt;
</code></pre>
<p>Infact, the original code seems to be executing fine if I remove the cast <code>.long()</code> on <code>window_elems</code>:</p>
<pre><code class="lang-auto"> window_elems = x[:, :, i:i + self.kernel, j:j + self.kernel].long()
print(window_elems.requires_grad,window_elems.grad_fn)
#False None
window_elems = x[:, :, i:i + self.kernel, j:j + self.kernel]
print(window_elems.requires_grad,window_elems.grad_fn)
#True &lt;SliceBackward0 object at 0x00000193CB7D80A0&gt;
</code></pre>

ptrblck: <p>Sorry, I meant <code>argmax</code> while you are using <code>amax</code>, which returns the <em>values</em> indeed.</p>