Bypass module during backpropagation

LuukJ: <p>Hi everyone,</p>
<p>I am currently working on a setup where the output of my network is first modified by my own module before computing the loss and backpropagating this, as follows:</p>
<pre><code>class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()

            self.Layer1 = Layer()
            self.Layer2 = Layer()

            self.Mod = Own_module()

            self.LossLayer = Loss()

    def forward(self, x):
        y1 = self.Layer1(x)
        y2 = self.Layer2(y1)

        self.output = self.Mod(y2)

    def backward(self):
        loss =  self.LossLayer(self.output)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
</code></pre>
<p>This works, however, <code>Own_mudule()</code> is an extensive simulation with lots of multiplications and loops (operations are completely fixed). Because of all the computational graphs that need to be saved during this simulation, memory issues occur. I tried solving this using:</p>
<pre><code>self.Mod.requires_grad_(False)
</code></pre>
<p>However, this still results in memory issues. And when I use <code>detach()</code> the backpropagation gets blocked completely.</p>
<p>Is there a way I can not save the computational graphs during the <code>Own_module()</code> and essentially bypass it during backpropagation?</p>
<p>Thanks a lot!</p>

KFrank: <p>Hi Luuk!</p>
<aside class="quote no-group" data-username="LuukJ" data-post="1" data-topic="139988" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/e56c9b/48.png" class="avatar"> LuukJ:</div>
<blockquote>
<p>Is there a way I can not save the computational graphs during the <code>Own_module()</code></p>
</blockquote>
</aside>
<p>Yes, but …</p>
<aside class="quote no-group">
<blockquote>
<p>and essentially bypass it during backpropagation?</p>
</blockquote>
</aside>
<p>You won’t want to bypass your <code>Own_module</code> backpropagation entirely.</p>
<p>Suppose <code>Own_module</code> were simply <code>return -input</code>, that is, it just flipped<br>
the sign of (the gradient of) the loss function?  Then you would be training<br>
your model to <em>maximize</em> the loss function, your predictions would become<br>
increasing poor, and your parameters would likely run off to infinity.</p>
<p>Therefore, you need to backpropagate back through <code>Own_module</code>  with<br>
at least some crude approximation to <code>Own_module</code>’s own gradient, at a<br>
minimum getting the sign of its gradient correct at least most of the time.</p>
<p>I would suggest that you package <code>Own_module</code> not as a <code>Module</code>, but<br>
as a <a href="https://pytorch.org/docs/1.10.1/autograd.html#torch.autograd.Function" rel="noopener nofollow ugc">torch.autograd.Function</a>.  Your custom <code>Function</code>’s <code>forward()</code><br>
method would implement your “extensive simulation” without saving<br>
the computational graph (using, e.g., <code>.detach()</code> or protected by a<br>
<code>with torch.no_grad():</code> block), and the computation graph would get<br>
“glued back together” because you would also implement your <code>Function</code>’s<br>
<code>.backward()</code> method.  The more accurately <code>.backward()</code> approximates<br>
the actual gradient of <code>.forward()</code>, the better your model will likely train.</p>
<p>As an aside, your <code>NeuralNet</code> <code>Module</code> doesn’t need to (and, to avoid<br>
confusion, shouldn’t) implement a <code>backward()</code> method – in a more<br>
standard design, the optimization code in your <code>.backward()</code> method<br>
would be moved outside of <code>NeuralNet</code> into your training loop after<br>
the application fo <code>NeuralNet</code> to your input data.</p>
<p>Best.</p>
<p>K. Frank</p>