Custom autograd function with helper functions

username: pylearn
cooked: <p>I defined a custom autograd function. It is very long. I want to convert the long computation into a few helper functions, so that it is more readable, and maybe the gradient/jacobian of the helper functions can be used directly?</p>
<p>The original implementation</p>
<pre><code class="lang-auto">class customFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, data, option):
        # some long computation
        res = ...
        # compute the gradient directly
        grad_input = ...     
        ctx.save_for_backward(grad_input)
        return res

    @staticmethod
    def backward(ctx, grad_output):
        grad_input, = ctx.saved_tensors
        return grad_output * grad_input, None, None
</code></pre>
<p>The converted code looks like</p>
<pre><code class="lang-auto">from torch.autograd.functional import jacobian

class customFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, data, option):
        # save data and option to be used in helper functions
        ctx.data = data
        ctx.option = option

        # get result with a helper function
        res1 = ctx.get_res1(x)
        res  = someRegularFunc(res1)

        # # compute the gradient by using the jacobian of `get_res1()`
        grad_res1_x = jacobian(get_res1, x)
        grad_input = ...
        ctx.save_for_backward(grad_input)
        return res

    def get_res1(ctx, x):
        data, option = ctx.data, ctx.option
        res1 = ...      # some short computation
        return res1  

    @staticmethod
    def backward(ctx, grad_output):
        grad_input, = ctx.saved_tensors
        return grad_output * grad_input, None, None
</code></pre>
<p>However, it does not work. Since <code>@staticmethod</code>cannot access the properties of the class itself</p>
<pre><code class="lang-auto">AttributeError: 'customFuncBackward' object has no attribute 'get_res1'
</code></pre>
<p>Is there a way to solve this problem?</p>

username: ptrblck
cooked: <p>Wouldn’t it work to define <code>get_res1</code> itself as a <code>staticmethod</code> or outside of the custom <code>autograd.Function</code>?</p>

username: pylearn
cooked: <p>Thanks for your quick reply.</p>
<ol>
<li>Defining <code>get_res1</code> as a <code>staticmethod</code> inside the custom <code>autograd.Function</code> does not work.</li>
</ol>
<p>(I am not familiar with python. Please correct me if I am wrong.)<br>
Reason: <code>staticmethod</code> can’t access or modify the class state. If <code>get_res1</code> is defined inside, then <code>forward</code> cannot use it.<br>
Reference: <a href="https://www.geeksforgeeks.org/class-method-vs-static-method-python/#:~:text=This%20method%20can%E2%80%99t%20access%20or%20modify%20the%20class%20state" class="inline-onebox" rel="noopener nofollow ugc">Class method vs Static method in Python - GeeksforGeeks</a>.</p>
<ol start="2">
<li>Defining <code>get_res1</code> as a regular function outside the custom <code>autograd.Function</code> works.</li>
</ol>
<p>Here is an example: to compute <code>y = 2(Ax + b)</code>, where <code>A=data</code> and <code>b=option</code></p>
<pre><code class="lang-auto">from torch.autograd.functional import jacobian

def get_res1(x, data, option):
    res1 = data@x + option     # some short computation
    return res1  

class customFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, data, option):
        # get result with a helper function
        res1 = get_res1(x, data, option)
        res  = 2 * res1

        # compute the gradient by using the jacobian of `get_res1()`
        grad_res1_x = jacobian(get_res1, (x, data, option))[0]
        grad_input = 2 * grad_res1_x
        ctx.save_for_backward(grad_input)
        return res

    @staticmethod
    def backward(ctx, grad_output):
        grad_input, = ctx.saved_tensors
        return grad_output * grad_input, None, None
</code></pre>
<p>Test</p>
<pre><code class="lang-auto">data = torch.tensor([[1., 2], [3, 4]])
option = torch.tensor([10., 10.])
x = torch.tensor([1., 1.], requires_grad=True)
s = customFunc.apply(x, data, option).sum()   # s = ((data@x + option)*2).sum()

s.backward()
x.grad           # -&gt; tensor([ 8., 12.])
</code></pre>
<p>However, it would be better if <code>get_res1</code> is bound to the custom <code>autograd.Function</code> since only the custom <code>autograd.Function</code> will use it.</p>
<p>Also, if it is defined outside as a regular function, then many inputs need to pass to <code>get_res1()</code> <strong>manually</strong>, while if it is a class method, then one can use <code>ctx.xxx</code> to access whatever properties/variables of the class.</p>
<p>My original code works (without any helper functions). I am just seeking a “better” way to implement, that is, to use helper function inside the class to increase readability (and to use jacobian directly in some case). However, if the helper functions need to be defined outside, it is also fine.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="pylearn" data-post="3" data-topic="181501">
<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/pylearn/48/54463_2.png" class="avatar"> pylearn:</div>
<blockquote>
<p>Reason: <code>staticmethod</code> can’t access or modify the class state.</p>
</blockquote>
</aside>
<p>I’m unsure which class state you are referring to as you would use the <code>ctx</code> object, wouldn’t you?</p>
<pre data-code-wrap="python"><code class="lang-python">class customFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, data, option):
        # save data and option to be used in helper functions
        ctx.data = data
        ctx.option = option

        # get result with a helper function
        res = customFunc.get_res1(ctx, x)
       
        return res

    @staticmethod
    def get_res1(ctx, x):
        data, option = ctx.data, ctx.option
        print(data, option)
        res1 = data
        return res1  

    @staticmethod
    def backward(ctx, grad_output):
        #grad_input, = ctx.saved_tensors
        return grad_output, None, None
    
fun = customFunc.apply
fun(torch.randn(1, 1), torch.rand(1, 1), torch.randn(1, 1))
</code></pre>

username: pylearn
cooked: <p>Now I see why I got the error. In my <code>forward</code>, I call the helper function by <code>ctx.get_res1()</code>, which returns error:</p>
<blockquote>
<p>AttributeError: ‘customFuncBackward’ object has no attribute ‘get_res1’</p>
</blockquote>
<p>Whereas in your <code>forward</code>, you use <code>customFunc.get_res1()</code>.</p>
<p>I thought the <code>ctx</code> is something like <code>self</code>, and we could use <code>self.some_method()</code>.</p>
<p>Thanks for your help.</p>

username: pylearn
cooked: <p>So is it true to say that <code>ctx</code> plays a different role as <code>self</code> in a regular class?</p>
<p>If so, it is fine to define or save some variables/attributes and use them later in other methods? That is, is it safe/correct do the following:</p>
<pre><code class="lang-auto">ctx.a = a   # in `forward` method
a = ctx.a   # in some other method
</code></pre>

username: ptrblck
cooked: <p>Yes, you can store arbitrary objects in the <code>ctx</code> context, but need to be careful with tensors.<br>
To store tensors you should use <code>ctx.save_for_backward</code> as you could create e.g. memory leaks otherwise.</p>