Using AD inside a custom defined torch.autograd.Function

James_Peraire: <p>Hi – so I’m defining a torch.autograd.Function.</p>
<pre><code class="lang-auto">class MyFunction(torch.autograd.Function):
  @staticmethod
  def forward(ctx, input1, input2):
    # some computation
    ctx.save_for_backward(....)
    return loss
  
  @staticmethod
  def backward(ctx, grad_output):
    # some computation
    return grad1, grad2
</code></pre>
<p>From what I understand, anything that happens to input1 inside of forward is not tracked and the gradient applied is whatever is returned in backward. (For example, if <code>loss = input1 * 2</code>, but <code>grad1 = grad_output * 10</code>, the computation graph won’t care about the first statement and treat the gradient as <code>grad_output * 10</code>). Is my understanding correct? Does torch treat the tensor differently because the operation happens inside a torch.autograd.Function rather than a module?</p>
<p>Assuming this understanding is correct, is there a way to combine both behaviors? In particular, suppose my forward function looks something like this:</p>
<pre><code class="lang-auto"># inside class MyFunction(torch.autograd.Function)
def forward(ctx, input1, input2):
  a = some_computation_a(input1)
  b = some_computation_b(input1)
  return a + b
</code></pre>
<p>All of the steps of <code>some_computation_a</code> aren’t directly differentiable so I want to supply the gradient to those operations myself, but all of the operations in <code>some_computation_b</code> are, so I’d like torch to track those operations automatically. Is this possible to do?</p>
<p>Thanks,<br>
James</p>

albanD: <p>Hi,</p>
<blockquote>
<p>Is my understanding correct?</p>
</blockquote>
<p>Yes you are absolutely correct!</p>
<blockquote>
<p>Does torch treat the tensor differently because the operation happens inside a torch.autograd.Function rather than a module?</p>
</blockquote>
<p>No, the Tensors are treated the same way, but everything runs with autograd disable diring the forward.</p>
<blockquote>
<p>Is this possible to do?</p>
</blockquote>
<p>Of course!<br>
I think the simplest solution is to (and is also a general guideline) keep your custom Function as small as possible: It will make the gradients easier to derive for you.<br>
In particular, you can have a custom function that handles fw/bw for <code>some_computation_a</code> only and then combine it is a regular python function:</p>
<pre data-code-wrap="python"><code class="lang-python">def my_fn(input1, input2):
    return SomeComputationA.apply(input1) + some_computation_b(input2)
</code></pre>
<p>Does that work for you?</p>

James_Peraire: <p>Yes I think so – but I’ll update this if for some reason it does not.</p>

111179: <p>Hi, Mr. albanD,</p>
<p>Can implement a learnable parameter by torch. autograd. Function? Like PRelu activation function implemented by torch. autograd. Function.</p>
<p>Thank you for your time</p>

albanD: <p>Hi,</p>
<p>I would recommend using a <code>nn.Module</code> for this. Where you can register your learnable parameters as <code>nn.Parameter</code> so that they are properly detected by other nn constructs.<br>
I would recommend using a custom autograd.Function only if you need to work with something that is non-differentiable or you want custom computations (beyond just gradient) to happen during the backward.</p>