Customizing torch.autograd.Function

username: Muhammad_Ali
cooked: <p>Hi there, hope all of you are fine.<br>
I am working on VQGAN+CLIP, and there they are doing this operation:</p>
<pre><code class="lang-auto">class ReplaceGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x_forward, x_backward):
        ctx.shape = x_backward.shape
        return x_forward

    @staticmethod
    def backward(ctx, grad_in):
        return None, grad_in.sum_to_size(ctx.shape)


class ClampWithGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, min, max):
        ctx.min = min
        ctx.max = max
        ctx.save_for_backward(input)
        return input.clamp(min, max)

    @staticmethod
    def backward(ctx, grad_in):
        (input,) = ctx.saved_tensors
        return (
            grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) &gt;= 0),
            None,
            None,
        )


replace_grad = ReplaceGrad.apply
clamp_with_grad = ClampWithGrad.apply

</code></pre>
<p>I am trying to understand this, but I cannot find any resources explaining it.<br>
I would be very grateful if anyone can help me with this. Like what is ReplaceGrad doing?</p>
<p>Thanks, for any help!</p>

username: albanD
cooked: <p>Hi,</p>
<p>A custom Function is simply a way to provide the backward formula directly.<br>
The Replace grad takes a and b as input and returns a. And in the backward, it will pass all the gradients toward b and no gradient toward a.</p>

username: Muhammad_Ali
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a> , hope you are fine.<br>
Thanks for helping me out.<br>
Sorry, I am not able to fully understand this. Can you share me some resources, where I can understand it…?</p>
<p>Thanks</p>

username: AlphaBetaGamma96
cooked: <p>The documentation has an example for a simple exponential function <a href="https://discuss.pytorch.org/t/how-to-customize-the-double-backward/17846/4">here</a></p>
<p>Another example from the forums is <a href="https://discuss.pytorch.org/t/how-to-customize-the-double-backward/17846/4">here</a>, although it’s for a custom Double Backward (so 2nd derivative) it might help illustrate what the custom <code>torch.autograd.Function</code> does and what you need code in for it to work.</p>

username: Muhammad_Ali
cooked: <p>Thanks, will look into it.</p>

username: Muhammad_Ali
cooked: <p>Sorry, I am not able to get the concept.<br>
I would be grateful, if you can explain me this piece of code:</p>
<pre><code class="lang-auto">class ClampWithGrad(torch.autograd.Function):
  @staticmethod
  def forward(ctx, input, min, max):
    ctx.min = min
    ctx.max = max
    ctx.save_for_backward(input)
    return input.clamp(min, max)

  @staticmethod
  def backward(ctx, grad_in):
    (input,) = ctx.saved_tensors
    return (
        grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) &gt;= 0),
        None, None,
    )
</code></pre>
<p>Like how is it trying to change grad_in, and why we have this ,None, None ?</p>
<p>Thanks</p>

username: AlphaBetaGamma96
cooked: <p>Ok, so after having a look on the forums there’s a link to how <code>torch.clamp</code> is implemented and it’s located <a href="https://github.com/pytorch/pytorch/blob/53fe804322640653d2dddaed394838b868ce9a26/torch/autograd/_functions/pointwise.py#L95" rel="noopener nofollow ugc">here</a>. And, I checked the gradient for that custom function and I’m pretty sure it’s wrong!</p>
<p>With regards to what <code>torch.autograd.Function</code> does, it’s a way (as <a class="mention" href="/u/alband">@albanD</a> said) to manually tell PyTorch what the derivative of a function should be (as opposed to getting the derivative automatically from Automatic Differentiation).</p>
<p>This version of the function is below</p>
<pre><code class="lang-auto">class ClampWithGradThatWorks(torch.autograd.Function):
  @staticmethod
  def forward(ctx, input, min, max):
    ctx.min = min
    ctx.max = max
    ctx.save_for_backward(input)
    return input.clamp(min, max)

  @staticmethod
  def backward(ctx, grad_out):
    input, = ctx.saved_tensors
    grad_in = grad_out*(input.ge(ctx.min) * input.le(ctx.max))
    return grad_in, None, None

</code></pre>
<p>I’ve renamed some variables to more appropriate names (like <code>grad_in</code> to <code>grad_out</code>).</p>
<p>The <code>forward</code> method of a function is that actual output the function returns, the <code>backward</code> method is where the gradients are calculated.</p>
<p>The <code>grad_out</code> term is the gradient of the Loss with respect to the <strong>output</strong> of the function, hence why I named it grad_<strong>out</strong>, and in order to get the gradient for the input terms (or any other terms within the function) we need to calculate the derivative of the output with respect to the other terms. So, in this case we need the gradient of the output with respect to</p>
<ol>
<li>input</li>
<li>min</li>
<li>max</li>
</ol>
<p>As the <code>min</code> and <code>max</code> are independent/fixed of the output their gradient is zero (which can be represented by <code>None</code>). So, that explains the <code>None, None</code> term in the output. But what about the <code> grad_in = grad_out*(input.ge(ctx.min) * input.le(ctx.max))</code> term?</p>
<p>Well, if you plot the clamp function out it’s effectively a linear function between the <code>min</code> and <code>max</code> and <code>min</code> for any input less than <code>min</code>, and <code>max</code> for any value greater than <code>max</code>. So, looking at the gradient it’s 0 for an <code>input</code> less than <code>min</code> and 0 for an <code>input</code> greater than <code>max</code>.</p>
<p>For any input in between <code>min</code> and <code>max</code> it’s 1. So, we need to define the gradient of this ‘in-between’ section as 1. This can be done by multiplying <code>input.ge(ctx.min)</code> with <code>input.le(ctx.max)</code>. This <em>also</em> defines the gradient of the previous 2 sections as well (gradient = 0 for <code>input &lt; min</code> and <code>input &gt; max</code>).</p>
<p>We’re <em>nearly done</em>, once we’ve define the gradient of the output with respect to the inputs, we need to connect the gradient of the loss with respect to the output (<code>grad_out</code>) in order to get the gradient of the loss with respect to the input which is <code>grad_in</code>. This can simply be done by using the <a href="https://en.wikipedia.org/wiki/Chain_rule" rel="noopener nofollow ugc">chain rule</a> which is defined by,</p>
<pre><code class="lang-auto">    grad_in = grad_out*(input.ge(ctx.min) * input.le(ctx.max))
</code></pre>
<p>Hopefully, this explains what <code>torch.autograd.Function</code> is doing! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Muhammad_Ali
cooked: <p>Thanks a loooot.<br>
It is so great of you, to explain in such detail.<br>
I will try out some examples, to get the most of it.</p>
<p>Thanks again…</p>

username: Muhammad_Ali
cooked: <p>ctx(in forward, and backward), is just acting like self here. Or does it has to do something special?</p>

username: AlphaBetaGamma96
cooked: <p>Kinda, except it has some methods of its own (more info <a href="https://discuss.pytorch.org/t/customizing-torch-autograd-function/143161/7">here</a>)</p>

username: Muhammad_Ali
cooked: <p>Hi,</p>
<pre><code class="lang-auto"> @staticmethod
  def backward(ctx, grad_out):
    input, = ctx.saved_tensors
    grad_in = grad_out*(input.ge(ctx.min) * input.le(ctx.max))
    return grad_in, None, None
</code></pre>
<p>Here grad_out, is what returned by forward?</p>

username: AlphaBetaGamma96
cooked: <p><code>grad_out</code> is returned via the autograd engine and is the gradient of the loss with respect to the output of your function, hence the name <code>grad_out</code>.</p>