Difference between apply an call for an autograd function

cdancette: <p>I created a simple autograd function, let’s call it F (based on torch.autograd.Function).</p>
<p>What’s the difference between calling</p>
<pre><code class="lang-auto">a = F.apply(args)
</code></pre>
<p>and instantiating, then calling, like this :</p>
<pre><code class="lang-auto">f = F()
a = f(args)
</code></pre>
<p>The two versions seem to be used in pytorch code, and in examples</p>

albanD: <p>Hi,</p>
<p>This difference is that instantiating + calling the <code>Function</code> works with “old style” functions (which are going to be deprecated in the future).<br>
Using <code>.apply</code> is for the “new style” functions. You can differentiate the two easily: new style functions are defined with only <code>@staticmethod</code>, while old style ones have an <code>__init__</code>.</p>

cdancette: <p>Ok, I see, thanks. Since I want to have a scalar parameter for my function, if I want to use the new style, I need to pass it as an argument to the forward then, and save it for backwards ?</p>
<p>I did this :</p>
<pre><code class="lang-auto">class F(torch.autograd.Function):
    def __init__(self, gamma=0.1):
        super().__init__()
        self.gamma = gamma

    def forward(self, args):
        pass
    def backward(self, args):
        pass
</code></pre>

albanD: <p>Hi,<br>
What you wrote is an old style function.</p>
<p>The equivalent new style would be:</p>
<pre><code class="lang-auto">class F_new(torch.autograd.Function):
    @staticmethod
    def forward(ctx, args, gamma):
        ctx.gamma = gamma
        pass

    @staticmethod
    def backward(ctx, args):
        pass

# Using your old style Function from your code sample:
F(gamma)(inp)
# Using the new style Function:
F_new.apply(inp, gamma)
</code></pre>

cdancette: <p>Thank you very much !</p>
<p>And so everytime I call F.apply, a new “instance” will be created with its own context ?</p>
<p>Can I save intermediary results with ctx.intermediary = intermediary ?</p>
<p>If I call several times the function before doing a backward pass, will it override the gamma attribute each time, and the intermediary results I saved with ctx ?</p>

albanD: <p>Each call to <code>.apply</code> will have a different context. So you can save everything you need in it without risk.<br>
Note that if you need to save input or outputs, you should use the <code>ctx.save_for_backward</code>.</p>

cdancette: <p>Ok, why do we need to use <code>save_for_backward</code> for this ? Is it just a convention ? Or does it performs additional checks ?<br>
But intermediary tensors are fine to save, right ? I tried to save some with <code>save_for_backwards</code> but it failed, so I saved them as attributes in self (ctx now)</p>

albanD: <p>Yes <code>save_for_backward</code> is just for input and outputs, it allows to perform additional checks (make sure that you don’t create non-collectable cycles). For intermediary results, you can save them as attribute of the context yes.</p>

cdancette: <p>Thanks, it much clear now !</p>

acgtyrant: <p>What version does the new style appear?</p>

albanD: <p>If I remember correctly, it was for <a href="https://github.com/pytorch/pytorch/releases/tag/v0.2.0">0.2.0</a> to support higher order derivatives.</p>

cosmozhang1988: <p>Hi Dear AlbanD, <a class="mention" href="/u/alband">@albanD</a></p>
<p>It seems the staticmethod is faster than the old sttyle function (ctx comparing with self).<br>
But what if I need to output some auxiliary output in the forward pass, which does not need to receive its gradient?<br>
In the old style, what I do is to create another instance method which is separate from the forward pass.<br>
A more concrete example:</p>
<pre><code class="lang-auto">class F(torch.autograd.Function):
    def __init__(self):
        super(F).__init__()

    def some_func(self, x2):
         //self.x is also needed here
         return z2

    def forward(self, x):
         self.x = x
         z1 = self.some_func(self.x)
         return y
                 
    def backward(self, y.grad):
         return x.grad

y = F()(x)
z2 = F()(x2)
y.backward()
</code></pre>
<p>Is there a way that this old styled can be re-rewriten into the new style (using ‘ctx’)?<br>
Thanks a lot!</p>

albanD: <p>Hi,</p>
<p>This is not what Function are made for. All arguments <em>must</em> be given as input to the forward function.<br>
To get this behavior, you can simply create an <code>nn.Module</code> which can do whatever you want.</p>

cosmozhang1988: <p>Yeah, my current solution is using <code>nn.module</code>, but comparing with <code>ctx</code>, it is much more slower…The reason I use <code>torch.autograd.Function</code> is because I need to write my own derivative.</p>

cosmozhang1988: <p><a class="mention" href="/u/alband">@albanD</a><br>
Maybe my description was not very precise. Taking <code>torch.autograd.Function</code> as an example, what I am seeking to do is:</p>
<pre><code class="lang-auto">class F(torch.autograd.Function):

   @staticmethod
    def some_func(z1, x2):
         //x2 is also needed here
         return z2

   @staticmethod
    def forward(ctx, x, x2):
         z1 = F.some_func3(x)
         y = z1**2
         z2 = F.some_func(z1, x)

         u = F.some_func2(y) \\derivative wrt x
         ctx.save_for_backward(u)
         return y, z2
                 
    @staticmethod
    def backward(ctx, y.grad, dummygrad):

         u, = ctx.saved_tensors
         return u, None
</code></pre>
<p>In this case, I do not need x2’s derivative, but I need to output z2 in the forward pass to save some computation cost in the intermediate values.</p>

cosmozhang1988: <p>Well, my way of writing actually seems working.</p>

s_n: <p>Hi <a class="mention" href="/u/alband">@albanD</a></p>
<p>Inside this forward function can I use register buffers?<br>
I can’t use self.register_buffer for sure. Also ctx.register_buffer doesn’t work.<br>
What is correct way to define register buffer in forward?</p>
<p>Thanks!</p>

albanD: <p>Hi,</p>
<p><code>.register_buffer()</code> is a function on nn.Modules, not autograd.Function.<br>
These are two completely different constructs. And you won’t need any buffer for custom Function as it is not a class instance, you should just create the Tensor you need during the forward.</p>

s_n: <p>Hi <a class="mention" href="/u/alband">@albanD</a> ,<br>
If I will just create a tensor how will that tensor be moved to GPU?<br>
I don’t want to use .to(device) because that won’t work with DataParallel.</p>
<p>Are you saying that all the tensors defined under autograd.Function will be automatically taken to device when I will do model.cuda() or torch.nn.DataParallel(model)?</p>
<p>Example:<br>
class Myclass(Function):<br>
<span class="mention">@staticmethod</span><br>
def forward(ctx,input):<br>
mytensor = torch.tensor([1])</p>
<p>So a model defined using above class (Myclass) when taken to device (by model.cuda() or torch.nn.DataParallel(model)) will take all the tensors in forward to the device?</p>

s_n: <p>I think I might have to create a tensor using register_buffer outside this forward function and then pass it to forward so that my tensors can be used with torch.nn.DataParallel and can be taken to device.</p>