Implementing custom backward function for grid_sample

username: eigenvivek
cooked: <p><strong>TLDR:</strong> is it possible to overwrite the gradient w.r.t. one of the inputs to a torch function?</p>
<p><code>torch.nn.functional.grid_sample</code> with <code>mode="nearest"</code> returns a gradient for <code>grid</code> that is zero everywhere (<a href="https://discuss.pytorch.org/t/does-the-nn-funtional-grid-sample-backwards-grad-to-grid/31736/4">as expected</a>).</p>
<p>I want overwrite this behavior and return a custom value for the gradient of <code>grid</code> (e.g., one everywhere or <code>None</code>). I also want to return the default gradient of <code>input</code>. I am trying to write a custom <code>torch.autograd.Function</code> to do this, but I am very confused by the syntax.</p>
<p>Here’s the default behavior:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

input = torch.randn(1, 1, 100, 101, 102, requires_grad=True)
grid = torch.randn(1, 1, 1, 200000, 3, requires_grad=True)

output = torch.nn.functional.grid_sample(
    input=input, grid=grid, mode="nearest", align_corners=False
)

y = output.sum()
y.backward()

input.grad.unique(), grid.grad.unique()
# prints : (tensor([0., 1., 2., 3., 4.]), tensor([0.]))
</code></pre>
<p>My interpretation of this is that the gradient for <code>input</code> is the number of times a voxel in <code>input</code> was sampled by <code>grid</code>; the gradient for <code>grid</code> is zero everywhere.</p>
<p>Following <a class="mention" href="/u/gfox">@gfox</a>’s <a href="https://discuss.pytorch.org/t/exluding-torch-clamp-from-backpropagation-as-tf-stop-gradient-in-tensorflow/52404/6">adaptation of <code>torch.clamp</code></a> to have a gradient of one everywhere, I tried to write a custom backward pass for <code>grid_sample</code> as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">from torch.cuda.amp import custom_bwd, custom_fwd


class DifferentiableNearestNeighbors(torch.autograd.Function):
    """
    In the forward pass, behaves like `torch.nn.functional.grid_sample`.
    In the backward pass, returns the default gradient for `input`, but
    also returns a gradient for `grid` that is 1 everywhere.
    """

    @staticmethod
    @custom_fwd
    def forward(ctx, input, grid, align_corners):
        ctx.save_for_backward(input, grid)
        return torch.nn.functional.grid_sample(
            input, grid, mode="nearest", align_corners=align_corners
        )

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        input, grid, = ctx.saved_tensors
        return grad_output.clone(), torch.zeros_like(grid), None


def differentiable_nearest_neighbors(input, grid, align_corners):
    return DifferentiableNearestNeighbors.apply(input, grid, align_corners)
</code></pre>
<p>However, <code>grad_output</code> is a tensor with the same shape as <code>output</code>, i.e., <code>(1, 1, 1, 1, 200000)</code>, but is one everywhere. My questions are then:</p>
<ul>
<li>How do you turn <code>grad_output</code> into the original gradient for <code>input</code> that <code>grid_sample</code> returns?</li>
<li>Is it possible to simply return the original gradient for <code>input</code> but somehow customize the gradient for <code>grid</code>?</li>
</ul>

username: KFrank
cooked: <p>Hi Vivek!</p>
<aside class="quote no-group" data-username="eigenvivek" data-post="1" data-topic="205694" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/eigenvivek/48/64833_2.png" class="avatar"> eigenvivek:</div>
<blockquote>
<p>I want overwrite this behavior and return a custom value for the gradient of <code>grid</code> (e.g., one everywhere or <code>None</code>). I also want to return the default gradient of <code>input</code>.</p>
</blockquote>
</aside>
<p>You can achieve this  by using <a href="https://pytorch.org/docs/2.3/generated/torch.Tensor.register_hook.html" rel="noopener nofollow ugc">register_hook()</a> on your <code>grid</code> tensor:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def grid_sample_with_hook (input, grid):
    _ = grid.register_hook (lambda grad: torch.ones_like (grad))
    return  torch.nn.functional.grid_sample (input, grid, mode = "nearest", align_corners = False)
    
input = torch.randn(1, 1, 100, 101, 102, requires_grad=True)
grid = torch.randn(1, 1, 1, 200000, 3, requires_grad=True)

print ('use grid_sample():')
output = torch.nn.functional.grid_sample(
    input=input, grid=grid, mode="nearest", align_corners=False
)

y = output.sum()
y.backward()

print ('input.grad.unique():', input.grad.unique())
print ('grid.grad.unique(): ', grid.grad.unique())

input.grad = None
grid.grad = None
print ('use grid_sample_with_hook():')
output = grid_sample_with_hook (input, grid)

y = output.sum()
y.backward()

print ('input.grad.unique():', input.grad.unique())
print ('grid.grad.unique(): ', grid.grad.unique())
</code></pre>
<p>yielding the following output:</p>
<pre><code class="lang-plaintext">2.3.1
use grid_sample():
input.grad.unique(): tensor([0., 1., 2., 3.])
grid.grad.unique():  tensor([0.])
use grid_sample_with_hook():
input.grad.unique(): tensor([0., 1., 2., 3.])
grid.grad.unique():  tensor([1.])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>I am trying to write a custom <code>torch.autograd.Function</code> to do this</p>
</blockquote>
</aside>
<p>I don’t think that this is a practical approach.  Once you put<br>
<code>grid_sample()</code> inside of an <code>autograd.Function</code>, it, itself, will no<br>
longer be tracked by autograd.  I suppose that you could somehow<br>
call <code>grid_sample()</code> twice, once to get autograd’s gradient for <code>input</code><br>
and then again inside of a <code>Function</code> so that autograd picks up your<br>
modified <code>grid.grad</code>, but doing so would seem round-about and inefficient.</p>
<p>Best.</p>
<p>K. Frank</p>