Nested custom autograd class

Jinghan_Yao: <p>Hi, I’d like to know how I can call the backward of a custom autograd module in the backward of another custom autograd module, when using activation checkpoint.</p>
<p>Following is an example:</p>
<pre><code class="lang-auto">import torch
from torch import autograd

class innerReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input &lt;= 0] = 0
        return grad_input

class OuterFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        with torch.enable_grad():
            relu_output = MyReLU.apply(input)
            this_output = relu_output * 2
            ctx.save_for_backward(input) # I do not want to save output, as it might take too much memory
            return this_output
    
    @staticmethod
    def backward(ctx, grad_output):
        input = ctx.saved_tensors
        # How to compute grad_input??
        return grad_input

</code></pre>
<p>One workaround I can think of is to write the backward process of the <code>innerRuLU</code> in <code>OuterFunction</code> backward function, but this is not feasible if the inner module is a complicated one.</p>
<p>Another way I tried was to call forward function and then use <code>autograd.grad</code> to compute the grad:</p>
<pre><code class="lang-auto">def backward(ctx, grad_output):
    input = ctx.saved_tensors
    relu_output = innerReLU.apply(input).requires_grad_()
    outer_result = relu_output * 2;
    grad_input = autograd.grad(outer_result, input, grad_output)
    return grad_input
</code></pre>
<p>But it says:</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>
<p>I would appreciate if anyone can shade some light on how to properly handle this.</p>
<p>Thanks!!</p>

ptrblck: <p>Autograd is disabled in custom <code>autograd.Function</code>s, so you might want to enable it in <code>innerReLU.forward</code>.</p>

soulitzer: <p>Are you trying to implement activation checkpoint yourself? Have you looked into <a href="https://pytorch.org/docs/stable/checkpoint.html" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.checkpoint — PyTorch 2.3 documentation</a>?</p>

Jinghan_Yao: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , thanks. But here I’d like to know if there are any ways to explicitly do the backward of <code>innerReLU</code> in the backward function of the <code>OuterFunction</code>. I do not necessarily need to use <code>autograd.grad</code>, just not sure of other ways to call the backward function of <code>innerReLU</code>.</p>

Jinghan_Yao: <p>Hi <a class="mention" href="/u/soulitzer">@soulitzer</a> , in this example code, yes, activation checkpoint is one thing I implement. But I am more curious on how to explicitly call backward of <code>innerReLU</code> in the backward function of <code>OuterFunction</code>.</p>
<p>Since the backward of <code>innerReLU</code> is defined, there should be some ways to call it, I am kinda interested in this.</p>

soulitzer: <p>Hmm may I ask why you need to activation checkpoint yourself?</p>
<p>Yes, it is possible to call it explicitly by accessing the attribute of the class.</p>
<pre><code class="lang-auto">class Test(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x
    @staticmethod
    def backward(ctx, x):
        return x * 2

a = torch.tensor(1., requires_grad=True)
out = Test.backward(None, a)
print(out)  # tensor(2., grad_fn=&lt;MulBackward0&gt;)
</code></pre>