torch.autograd.Function's backward is not triggered when func.apply() is used

username: puranjay_mishra
cooked: <p>Hello. I’m trying to use torch.autograd.Function to define a function with a custom forward and backward pass. I’m following the guide in the pytorch documentation which gives the following example -</p>
<pre><code class="lang-auto">&gt;&gt;&gt; class Exp(Function):
&gt;&gt;&gt;     @staticmethod
&gt;&gt;&gt;     def forward(ctx, i):
&gt;&gt;&gt;         result = i.exp()
&gt;&gt;&gt;         ctx.save_for_backward(result)
&gt;&gt;&gt;         return result
&gt;&gt;&gt;
&gt;&gt;&gt;     @staticmethod
&gt;&gt;&gt;     def backward(ctx, grad_output):
&gt;&gt;&gt;         result, = ctx.saved_tensors
&gt;&gt;&gt;         return grad_output * result
&gt;&gt;&gt;
&gt;&gt;&gt; # Use it by calling the apply method:
&gt;&gt;&gt; output = Exp.apply(input)
</code></pre>
<p>my code looks something like this -</p>
<pre><code class="lang-auto">class Exp(Function):
     @staticmethod
     def forward(ctx, i,):
          #I do some changes to the input i which works as intended
          return i
     @staticmethod
     def backward(ctx,grad_output):
          grad_output = torch.relu(grad_output)
          print(grad_output.shape)
          return grad_output
</code></pre>
<p>I expected this to print grad_output shape everytime the function has been called. This function has called in the forward method in the torch.nn.Module subclass (which looks really similar to the Linear layer). Let’s called this class Cats. Cats has been instantiated in the init method of a model.<br>
But the code does not print the grad_output. What could be the reason?</p>

username: AlphaBetaGamma96
cooked: <p>Have you called <code>.backward()</code> anywhere in your code?</p>
<p>Also, can you post a minimal reproducible example to replicate the error?</p>

username: puranjay_mishra
cooked: <p>Hello! Thanks for the reply.<br>
I did not explicitly call the <code>.backward()</code> , however i did apply it according to the pytorch documentation by typing <code>Exp.apply(input)</code><br>
The following is the code for reproducing the error(provided the proper imports have made) -</p>
<pre><code class="lang-auto">class autograd_func(torch.autograd.Function):
    def forward(self,layer):
        alpha = torch.ones(layer.shape[0])
        Alpha = (alpha)
        return layer,Alpha
    def backward(self,grad_output):
        print("checking backward call for autograd_func")
        return grad_output


class custom_linear(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = False,
                 device=None, dtype=None):
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(custom_linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))
        if bias:
            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self) -&gt; None:
        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
        # https://github.com/pytorch/pytorch/issues/57109
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in &gt; 0 else 0
            init.uniform_(self.bias, -bound, bound)

    def forward(self,input):
        A,B = autograd_func.apply(self.weight.data)
        self.weight.data = A
        #B does get used in the complete implementation but it does not make a difference
        out = F.linear(input, self.weight, self.bias)
        return out

    def extra_repr(self) -&gt; str:
        return 'in_features={}, out_features={}, bias={}'.format(
            self.in_features, self.out_features, self.bias is not None)


x = torch.randn([1,100])
y_star = torch.ones([1,10])
cust_layer = custom_linear(100,10)
y = cust_layer(x)
# print(y.shape,y_star.shape)
criterion = torch.nn.MSELoss()
loss = criterion(y,y_star)
loss.backward()
</code></pre>
<p>I hope this helps!</p>

username: AlphaBetaGamma96
cooked: <aside class="quote no-group" data-username="puranjay_mishra" data-post="3" data-topic="147673">
<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/p/b5ac83/48.png" class="avatar"> puranjay_mishra:</div>
<blockquote>
<pre><code class="lang-auto">    def forward(self,layer):
</code></pre>
</blockquote>
</aside>
<p>the forward and backward methods need to be decorated with <code>@staticmethod</code>, add that above each method and see if that fixes the problem</p>

username: puranjay_mishra
cooked: <p>I did that, did not work…</p>

username: AlphaBetaGamma96
cooked: <p>Ok, so the comments about <code>@staticmethod</code> were correct, it’s just that you need to add a few more things. For example change <code>self</code> to <code>ctx</code>, and remove the calls to <code>.data</code>, for example <code>A,B = autograd_func.apply(self.weight.data)</code> as <code>.data</code> is deprecated.  Make sure to follow the exact layout of the arguments for the <code>forward</code> and <code>backward</code> in this example <a href="https://pytorch.org/docs/stable/notes/extending.html#example" rel="noopener nofollow ugc">below</a>. Do make sure to read through the tutorial (<a href="https://pytorch.org/docs/stable/notes/extending.html" class="inline-onebox" rel="noopener nofollow ugc">Extending PyTorch — PyTorch 1.11.0 documentation</a>) too.</p>
<pre><code class="lang-auto"># Inherit from Function
class LinearFunction(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    # bias is an optional argument
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias
</code></pre>

username: InnovArul
cooked: <p>This simple code works for me.</p>
<pre><code class="lang-auto">import torch

class autograd_func(torch.autograd.Function):
    @staticmethod
    def forward(self,layer):
        alpha = torch.ones(layer.shape[0])
        Alpha = (alpha)
        return layer,Alpha
    
    @staticmethod
    def backward(self,grad_output, other):
        print("checking backward call for autograd_func")
        return grad_output


if __name__ == '__main__':
    x = torch.randn(5,6)
    x.requires_grad = True
    y, a = autograd_func.apply(x)
    y.sum().backward()
</code></pre>

username: puranjay_mishra
cooked: <p>Thanks for the reply and the links!<br>
Changing it to the ctx did not work.I was unaware of the fact that .data is depreceated. I will rewrite that part of code and try it again, as autograd does not track changed by making a backwards graph when .data is used(correct me if I’m wrong here)<br>
I will try out your suggestions and update you.</p>

username: puranjay_mishra
cooked: <p>Thank you for the reply!<br>
This did work, yes. The torch.randn is a tensor whose computations are tracked by autograd. Using backward on .data is depreceated and changes on it are not tracked by it. Ill make the necessary changes and update you.</p>

username: puranjay_mishra
cooked: <p>It did work! Thank you for your feedback InnovArul and AlphaBetaGamma96</p>