When should you save_for_backward vs. storing in ctx

Jason_Hirshman: <p>I’m defining a new function using the 0.2 style and am wondering when it is appropriate to store intermediate results in the ctx object as opposed to using the save_for_backward function.</p>
<p>Here is a simple example:</p>
<pre><code class="lang-python"># OPTION 1
class Square(Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        c = a + b
        return c * c

    @staticmethod
    def backward(ctx, grad_output):
        a, b = ctx.saved_variables
        c = a + b

        grad = grad_output * (2 * c)
        return grad, grad


# OPTION 2
class Square(Function):
    @staticmethod
    def forward(ctx, a, b):
        c = a + b
        ctx.intermediate_results = c
        return c * c

    @staticmethod
    def backward(ctx, grad_output):
        c = ctx.intermediate_results

        grad = grad_output * (2 * c)
        return grad, grad

</code></pre>
<p>What are the advantages/disadvantages to Option 1 over Option 2?  When would you want to use one over the other?</p>
<p>(In a real example, the computation of the intermediate variable “c” would be a lot more intensive, e.g. involve O(n^2) operations)</p>
<p>Thank you.</p>

tom: <p>I think you can also store results.<br>
My impression is that it is “use <code>save_for_backward</code> if there isn’t a reason not to”. Efficiency is a reason.</p>
<p>Best regards</p>
<p>Thomas</p>

Jason_Hirshman: <p>Ok, but is there a reason why you should use save_for_backward?  How much more efficient does it have to be to justify storing directly to ctx?</p>

tom: <p>I don’t think there are hard rules about it, so I would look at the torch.autograd._functions for inspiration when to use which and see what looks similar to your case.<br>
For example in the <a href="https://github.com/pytorch/pytorch/blob/master/torch/autograd/_functions/pointwise.py#L142" rel="nofollow noopener">Clamp Function</a> you store the mask in the <code>ctx</code> while for the <a href="https://github.com/pytorch/pytorch/blob/master/torch/autograd/_functions/pointwise.py#L155" rel="nofollow noopener">square root</a> a few lines down, you just store the input.</p>
<p>Best regards</p>
<p>Thomas</p>

Jason_Hirshman: <p>Thanks, Thomas.  Looking through the source code it seems like the main advantage to save_for_backward is that the saving is done in C rather python.   So it seems like anytime the computation for the intermediate result would allocate python storage space anyway, it makes sense to use the ctx object.</p>

pedrodiamel: <p>I think that option 1 delete the gpu memory of the Variable.</p>

jxgu1016: <p>there will be memory leak in option 2, where ctx.intermediate_results won’t be released.</p>

tom: <p>Hu? The ctx would go out of scope after the backward, too, would it not?</p>
<p>Best regards</p>
<p>Thomas</p>

aviolov: <p>We experienced a memory leak on a custom module using custom-make GPU kernels. Storing intermediate results in <code>ctx.x, ctx.z</code> etc resulted in them not being released over multiple mini-batches, quickly exhausting GPU memory. We did not experience this in pytorch 0.3.1, it sounds like it is due to some changes in pytorch 0.4. Either explicitly <code>del ctx.x</code> etc or adding x,z to <code>.save_for_backward</code> both fixed the memory leak. I would still like to get a good explanation for why you would want to use <code>save_for_backward</code> vs. directly assigning to <code>ctx</code>.</p>
<p>I seem to recall that in earlier torch versions it was not possible to <code>save_for_backward</code> an intermediate result, only inputs to the <code>forward</code> function, but that seems to be no longer the case.</p>

eduardo4jesus: <p>The fact that <code>save_for_backward</code> can only save variables is, in my application, a good reason to use <em>store in <code>ctx</code></em>.</p>

Connor_Holmes: <p>Not sure if I am too late to reply, but there is some information about this at this <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">link</a>, under “How to use: Step 2”. It seems like anything that is a tensor should be stored using <code>save_for_backward</code>, but anything else should be stored directly on <code>ctx</code>.</p>

soulitzer: <p>save_for_backward gives you two benefits:</p>
<ol>
<li>if the output is saved for backward, runs special logic necessary to make sure that this does not result in a reference cycle</li>
<li>it allows the saved tensor to be cleared immediately after custom backward is finished running</li>
</ol>
<p>If a tensor is saved directly on ctx, it will be kept alive for the duration of the autograd graph, which normally is OK if you make sure it is cleared. Using save_for_backward in this situation allows it to be cleared sooner though.</p>