Legacy autograd function with non-static forward method is deprecated. Please use new-style autograd function with static forward method

HamidGadirov: <p>Hi, I am getting error messages after trying to update the neural network (for optical flow) for the current PyTorch version (1.13.0+cu117, previous implementation was for 1.1). The code is the following:</p>
<pre><code class="lang-auto">class CorrelationFunction(Function):
    def __init__(self, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        super(CorrelationFunction, self).__init__()
        print("in CorrelationFunction __init__")
        self.pad_size = pad_size
        self.kernel_size = kernel_size
        self.max_displacement = max_displacement
        self.stride1 = stride1
        self.stride2 = stride2
        self.corr_multiply = corr_multiply
        # self.out_channel = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1)

    # @staticmethod
    def forward(self, input1, input2): # self,
        print("in CorrelationFunction __init__")
        self.save_for_backward(input1, input2)

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()
            output = input1.new()

            correlation_cuda.forward(input1, input2, rbot1, rbot2, output, 
                self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)

        print("returning CorrelationFunction forward")
        return output

    # @staticmethod
    def backward(self, grad_output): # self, 
        input1, input2 = self.saved_tensors

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()

            grad_input1 = input1.new()
            grad_input2 = input2.new()

            correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,
                self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)

        return grad_input1, grad_input2


class Correlation(Module):
    def __init__(self, pad_size=0, kernel_size=0, max_displacement=0, stride1=1, stride2=2, corr_multiply=1):
        super(Correlation, self).__init__()
        print("in Correlation __init__")
        self.pad_size = pad_size
        self.kernel_size = kernel_size
        self.max_displacement = max_displacement
        self.stride1 = stride1
        self.stride2 = stride2
        self.corr_multiply = corr_multiply


    # @staticmethod
    def forward(self, input1, input2):
        print("in Correlation forward")

        result = CorrelationFunction(self.pad_size, self.kernel_size, self.max_displacement, self.stride1, self.stride2, self.corr_multiply)(input1, input2)

        return result
</code></pre>
<p>and I am calling <code>Correlation</code> like this:<br>
<code>out_corr_1 = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(feature_1, feature_2_warp)</code></p>
<p>However, there is an error:<br>
<code>RuntimeError: Legacy autograd function with non-static forward method is deprecated. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)</code></p>
<p>I have seen a related discussion (<a href="https://discuss.pytorch.org/t/legacy-autograd-function-with-non-static-forward-method-is-deprecated-and-will-be-removed-in-1-3/69869/4">here</a>), where it was suggested to “remove the comments before <code>@staticmethod</code> (<code>#</code> )” and then use “Class_name.apply(x)” by <a class="mention" href="/u/ptrblck">@ptrblck</a>, but then I have another issue:<br>
<code>TypeError: apply() takes 2 positional arguments but 3 were given</code></p>
<p>Thanks in advance!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/hamidgadirov">@HamidGadirov</a>,</p>
<p>Have you read the example in <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.1 documentation</a> ?</p>
<p>The signature for <code>torch.autograd.Function</code> is now defined as,</p>
<pre><code class="lang-auto">class CustomFunc(torch.autograd.Function)

  #no need for the constuctor (__init__) bit 

  @staticmethod
  def forward(ctx, *args, **kwargs):
     #forward pass here

  @staticmethod
  def backward(ctx, *args, **kwargs):
    #backward pass here
</code></pre>
<p>You need to remove the <code>__init__</code> method from the class (like I showed above), and then when you want to use the Function it needs to be defined as,</p>
<pre><code class="lang-auto">func = CustomFunc.apply(*args, **kwargs)
</code></pre>
<p>Now, <code>func</code> is a callable instance of <code>torch.autograd.Function</code>.</p>
<aside class="quote no-group" data-username="HamidGadirov" data-post="1" data-topic="165375">
<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/h/eb9ed0/48.png" class="avatar"> HamidGadirov:</div>
<blockquote>
<p><code>TypeError: apply() takes 2 positional arguments but 3 were given</code></p>
</blockquote>
</aside>
<p>From this error message, check you’re not passing the <code>self</code> arg to the <code>.apply()</code> call. Also, I think that should be defined as <code>ctx</code> instead of <code>self</code>, but it could be arbitrary. See if it makes a difference.</p>

HamidGadirov: <p>Thanks <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> for reply, I reimplemented the code for PyTorch 1.13:</p>
<pre><code class="lang-auto">
class CorrelationFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input1, input2, 
            pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        print("in CorrelationFunction forward")
        ctx.save_for_backward(input1, input2)

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()
            output = input1.new()

            correlation_cuda.forward(input1, input2, rbot1, rbot2, output, 
                pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input1, input2 = ctx.saved_tensors

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()

            grad_input1 = input1.new()
            grad_input2 = input2.new()

            correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,
                pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)

        return grad_input1, grad_input2

class Correlation(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input1, input2, 
            pad_size=0, kernel_size=0, max_displacement=0, stride1=1, stride2=2, corr_multiply=1): 
        print("in Correlation forward")

        result = CorrelationFunction(input1, input2, pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)
        # result = CorrelationFunction(pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)(input1, input2)

        return result
</code></pre>
<p>However, now when I try to use the output of <code>Correlation</code>, e.g.:</p>
<p><code>out_corr_relu_1 = self.leakyRELU(out_corr_1)</code></p>
<p>there is an error:</p>
<p><code>TypeError: leaky_relu_(): argument 'input' (position 1) must be Tensor, not CorrelationFunction</code></p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/hamidgadirov">@HamidGadirov</a>,</p>
<p>You need to check wherever you’ve defined a <code>self.leakyRELU(out_corr_1)</code>. Because it seems that <code>out_corr_1</code> is your custom autograd Function and not a Tensor. I assume this is within your <code>correlation_cuda</code> function?</p>

HamidGadirov: <p>Even if I comment out <code>correlation_cuda</code> function, it is the same issue. The problem is that return type of <code>out_corr_1</code> is CorrelationFunction, and I don’t know why. Converting to tensor doesn’t work either.</p>

AlphaBetaGamma96: <p>Can you share a minimal reproducible example? It seems your using the custom function somewhere else. Did you remember to use the <code>.apply</code> method?</p>
<p>You can do something like this,</p>
<pre><code class="lang-auto">def func(*args, **kwargs):
  return CustomFunc.apply(*args, **kwargs)

func(*args, **kwargs) #call function like a normal python function.
</code></pre>

HamidGadirov: <p>I was using .apply() of course. here is the minimal reproducible example (combination of pytorch example and my code):</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square, you can specify with a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class CorrelationFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input1, input2, 
            pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        print("in CorrelationFunction forward")
        ctx.save_for_backward(input1, input2)

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()
            output = input1.new()

            # correlation_cuda.forward(input1, input2, rbot1, rbot2, output, 
            #     pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input1, input2 = ctx.saved_tensors

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()

            grad_input1 = input1.new()
            grad_input2 = input2.new()

            # correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,
            #     pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)

        return grad_input1, grad_input2

class Correlation(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input1, input2, 
            pad_size=0, kernel_size=0, max_displacement=0, stride1=1, stride2=2, corr_multiply=1): 
            # self, self for the object; when you call, it is without - synthactic sugar!
        print("in Correlation forward")

        result = CorrelationFunction(input1, input2, pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)
        # result = CorrelationFunction(pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)(input1, input2)

        return result

net = Net()
input = torch.randn(1, 1, 32, 32)

# test
leakyRELU = nn.LeakyReLU(0.1, inplace=True)
out_corr_1 = Correlation.apply(input, input, 4, 1, 4, 1, 1, 1)
out_corr_relu_1 = leakyRELU(out_corr_1)
</code></pre>
<p>it throws the same <code>must be Tensor, not CorrelationFunction</code> error</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="HamidGadirov" data-post="7" data-topic="165375">
<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/h/eb9ed0/48.png" class="avatar"> HamidGadirov:</div>
<blockquote>
<pre><code class="lang-auto">        result = CorrelationFunction(input1, input2, pad_size, kernel_size, max_displacement, stride1, stride2, corr_multiply)
</code></pre>
</blockquote>
</aside>
<p>You need to use <code>.apply</code> here too as <code>CorrelationFunction</code> is <em>also</em> a <code>torch.autograd.Function</code> object. That’s why you’re getting your error about it not being a Tensor.</p>

HamidGadirov: <p>Thanks, he issue is solved now.</p>