Custom CUDA extension backward pass does not appear as grad_fn

username: mesh-god
cooked: <p>I am implementing a basic sharpener with one parameter just to test the CUDA extension and get it working for further work. The python class is defined as so.</p>
<pre><code class="lang-auto">class NASFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, A):
        output = nas_cuda.forward(input, A)
        variables = output
        ctx.save_for_backward(*variables)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        output = nas_cuda.backward(grad_output[0], ctx.saved_variables[1])
        return output, ctx.saved_variables[1]


class NAS(nn.Module):
        def __init__(self):
            super(NAS, self).__init__()
            self.A = torch.nn.Parameter(torch.ones(1) * 0.5).cuda()

        def forward(self, input):
            return NASFunction.apply(input, self.A)
</code></pre>
<p>I then define a forward pass and a dummy backward pass. I cannot invoke the dummy backward pass in order to test it and implement it further as it doesn’t appear as a grad_fn hence it can never be applied. The cuda implementation is below.</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

#include &lt;cuda.h&gt;
#include &lt;cuda_runtime.h&gt;

#include &lt;vector&gt;


template &lt;typename scalar_t&gt;
__global__ void nas_cuda_forward_kernel(
    torch::PackedTensorAccessor&lt;scalar_t, 4, torch::RestrictPtrTraits, size_t&gt; input,
    torch::PackedTensorAccessor&lt;scalar_t, 4, torch::RestrictPtrTraits, size_t&gt; output,
    torch::PackedTensorAccessor&lt;scalar_t, 1, torch::RestrictPtrTraits, size_t&gt; Ap
){

  int i = blockIdx.x * blockDim.x + threadIdx.x;
  int j = blockIdx.y * blockDim.y + threadIdx.y;

  if (i &lt; 1) i = 1;
  if (i &gt; 254) i = 254;
  if (j &lt; 1) j = 1;
  if (j &gt; 254) j = 254;

  // a b c   0 A 0 
  // d e f . A 1 A
  // g h i   0 A 0 

  float norm_factor = (1 + 4 * 0.5);

  const auto b_r = input[0][0][i-1][j];
  const auto b_g = input[0][1][i-1][j];
  const auto b_b = input[0][2][i-1][j];

  const auto d_r = input[0][0][i][j-1];
  const auto d_g = input[0][1][i][j-1];
  const auto d_b = input[0][2][i][j-1];
  const auto e_r = input[0][0][i][j];
  const auto e_g = input[0][1][i][j];
  const auto e_b = input[0][2][i][j];
  const auto f_r = input[0][0][i][j+1];
  const auto f_g = input[0][1][i][j+1];
  const auto f_b = input[0][2][i][j+1];

  const auto h_r = input[0][0][i+1][j];
  const auto h_g = input[0][1][i+1][j];
  const auto h_b = input[0][2][i+1][j];

  float Ap_ = Ap[0];

  output[0][0][i][j] = (e_r + b_r * Ap_ + d_r * Ap_ + f_r * Ap_ + h_r * Ap_) / norm_factor;
  output[0][1][i][j] = (e_g + b_g * Ap_ + d_g * Ap_ + f_g * Ap_ + h_g * Ap_) / norm_factor;
  output[0][2][i][j] = (e_b + b_b * Ap_ + d_b * Ap_ + f_b * Ap_ + h_b * Ap_) / norm_factor;

}

std::vector&lt;torch::Tensor&gt; nas_cuda_forward(
  torch::Tensor input,
  torch::Tensor A
){
  auto output = torch::zeros_like(input);
  const auto batch_size = input.size(0);
  const auto channels = input.size(1);
  const auto h = input.size(2);
  const auto w = input.size(3);

  const dim3 blockDim(16, 16);
  const dim3 gridDim(16, 16);

  AT_DISPATCH_FLOATING_TYPES(input.type(), "nas_cuda_forward", ([&amp;] {
    nas_cuda_forward_kernel&lt;scalar_t&gt;&lt;&lt;&lt;gridDim, blockDim&gt;&gt;&gt;(
      input.packed_accessor&lt;scalar_t, 4, torch::RestrictPtrTraits, size_t&gt;(),
      output.packed_accessor&lt;scalar_t, 4, torch::RestrictPtrTraits, size_t&gt;(),
      A.packed_accessor&lt;scalar_t, 1, torch::RestrictPtrTraits, size_t&gt;()
      );
  }));

  return {output, A};
}


std::vector&lt;torch::Tensor&gt; nas_cuda_backward(
  torch::Tensor grad_input,
  torch::Tensor A
){
    A = A* 0.5;
    return {grad_input, A};
}
</code></pre>
<p>Then if I execute the following python to test it</p>
<pre><code class="lang-auto">import torch

from cuda.nas import NAS

x = torch.ones(1,3,256,256, requires_grad=True).cuda()

nas = NAS()

y = nas(x)
</code></pre>
<p>Which provides the following output</p>
<pre><code class="lang-auto">         [[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          ...,
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]],

         [[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          ...,
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 1., 1.,  ..., 1., 1., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]]], device='cuda:0')
</code></pre>
<p>And the A parameter is returned as</p>
<pre><code class="lang-auto">tensor([0.5000], device='cuda:0', grad_fn=&lt;ToCopyBackward0&gt;)
</code></pre>
<p>There is no grad_fn in the output tensor x hence i cannot use Autograd to propagate any information back.</p>
<p>What am I missing here? I am failing to see why output of nas has no grad_fn</p>

username: ptrblck
cooked: <p>You are creating non-leaf tensors by calling differentiable operations on the parameters:</p>
<pre><code class="lang-python">x = torch.ones(1,3,256,256, requires_grad=True).cuda()
...
self.A = torch.nn.Parameter(torch.ones(1) * 0.5).cuda()
</code></pre>
<p>Call the <code>cuda()</code> operation (and any other <code>to()</code> operation) on the tensor before wrapping it into an <code>nn.Parameter</code> or set it as an attribute directly:</p>
<pre><code class="lang-python">x = torch.ones(1,3,256,256, requires_grad=True, device="cuda")
...
self.A = torch.nn.Parameter(torch.ones(1, device="cuda") * 0.5)
</code></pre>

username: mesh-god
cooked: <p>Hi, thanks so much for the guidance.</p>
<p>I implemented you’re changes and the gradient requirement disappears</p>
<pre><code class="lang-auto">x = torch.ones(1,3,256,256, requires_grad=True, device="cuda")
</code></pre>
<pre><code class="lang-auto">(Pdb) x.requires_grad
True
</code></pre>
<p>which when passed through the model shows</p>
<pre><code class="lang-auto">(Pdb) nas(x)[0].requires_grad
False
</code></pre>
<p>Is it perhaps an issue with the my cuda kernel and it prevents it from being able to track gradient information?</p>
<p>Thanks</p>

username: ptrblck
cooked: <p>Could you create a GitHub repository containing a minimal executable code snippet as well as the build instructions to reproduce the issue, please?<br>
Right now it seems the binding from your extension is missing to the Python API.</p>

username: mesh-god
cooked: <p>Sure, thank you, will share this with you imminently</p>