Implementing a custom convolution using conv2d_input and conv2d_weight

fsds: <p>Hi,<br>
I have been trying to implement a custom convolutional layer.<br>
In order to do that, I’m using torch.nn.functional.conv2d in the forward pass, and both torch.nn.grad.conv2d_weight and torch.nn.grad.conv2d_input in the backward pass.<br>
I started getting OOM exceptions when entering torch.nn.grad.conv2d_weight.</p>
<p>My question is, what exactly is the difference between using:</p>
<pre><code class="lang-auto">torch.nn.functional.conv2d(x, w)
</code></pre>
<p>and</p>
<pre><code class="lang-auto">MyConv().apply(x, w)
</code></pre>
<p>when MyConv is implemented as follows:</p>
<pre><code class="lang-auto">class MyConv(Function):
   @staticmethod
   def forward(ctx, x, w):
     ctx.save_for_backward(x, w)
     return F.conv2d(x, w)

  @staticmethod
  def backward(ctx, grad_output):
    x, w = ctx.saved_variables
    x_grad = w_grad = None
    if ctx.needs_input_grad[0]:
      x_grad = torch.nn.grad.conv2d_input(x.shape, w, grad_output)
    if ctx.needs_input_grad[1]:
      w_grad = torch.nn.grad.conv2d_weight(x, w.shape, grad_output)
    return x_grad, w_grad
</code></pre>
<p>Why would torch.nn.grad.conv2d_weight return an OOM exception when torch.nn.functional.conv2d (that I assume also uses torch.nn.grad.conv2d_weight in the backward pass) did not?</p>
<p>Thanks.</p>

Tony_Lee: <p>hi , have you solved your problem ? I want to define a conv2d layer too, can you share me you code?</p>

fsds: <p>Yes.<br>
I’ve avoided this by directly calling cudnn_convolution_backward_input and cudnn_convolution_backward_weight (by following this example <a href="https://github.com/pytorch/extension-cpp" rel="nofollow noopener">https://github.com/pytorch/extension-cpp</a>, and adding two c++ functions that call the cudnn functions) instead of using torch.nn.grad.conv2d_input and torch.nn.grad.conv2d_weight.<br>
You might not have that problem though (depends on your nn and GPU model). You should first check if the if torch.nn.grad.conv2d_weight and torch.nn.grad.conv2d_input are working for your model without returning an out of memory exception.</p>

Tony_Lee: <p>I use your code above to modify my code, but I got an error.<br>
Is F.conv2d(x, w) have both forward and backward methords?</p>

Tony_Lee: <p>I just want to custom a similar conv2d layer as the API<br>
Can I just modify the pytorch python interface?</p>

fsds: <p>What exactly is the error?</p>

Tony_Lee: <p>it shows as follows</p>
<p>File “/home/lth/anaconda3/lib/python3.5/site-packages/torch/nn/functional.py”, line 90, in conv2d<br>
return f(input, weight, bias)</p>
<p>TypeError: argument 0 is not a Variable</p>
<p>my code is like that:<br>
class Conv2d(_ConvNd):</p>
<pre><code>def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
    kernel_size = _pair(kernel_size)
    stride = _pair(stride)
    padding = _pair(padding)
    dilation = _pair(dilation)
    super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False, _pair(0), groups, bias)

def forward(self, input):       
    return conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
</code></pre>
<p>conv2d = Conv2dF.apply</p>
<p>class Conv2dF(Function):</p>
<pre><code>@staticmethod
def forward(cxt, input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
        
    cxt.save_for_backward(input, weight, bias)

    return F.conv2d(input, weight, bias, stride, padding, dilation, groups)


@staticmethod
def backward(cxt, grad_output):
    input, weight, bias = cxt.saved_variables
            
    grad_input = grad_weight= grad_bias = None

    if cxt.needs_input_grad[0]:
        grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output)
        
    if cxt.needs_input_grad[1]:
        grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)
            
    if bias is not None and cxt.needs_input_grad[2]:
        grad_bias = grad_output.sum(0).squeeze(0)
    
    if bias is not None:
        return grad_input, grad_weight, grad_bias
    else:
        return grad_input, grad_weight
</code></pre>
<p>Thanks so much !</p>

fsds: <p>It looks like you’re using an old version of pytorch. Try moving to pytorch 0.4.0 and see if it works.<br>
To verify that use:</p>
<pre><code class="lang-auto">import torch
print(torch.__version__) #should be 0.4.0
</code></pre>

Tony_Lee: <p>I have just update the pytorch version to 0.4.0,but it  is also that error<br>
could you show me your demo code<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=5" title=":sob:" class="emoji" alt=":sob:"><img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=5" title=":sob:" class="emoji" alt=":sob:"><img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=5" title=":sob:" class="emoji" alt=":sob:"></p>

negar_goli: <p>Could you please share your code that you are calling cudnn_convolution_weight</p>

hanspinckaers: <p>For other people googling this, I posted some code in this thread: <a href="https://discuss.pytorch.org/t/cuda-error-with-cudnn-convolution-backward-weight-function/41214/2" class="inline-onebox">Cuda error with cudnn convolution backward weight function</a></p>

Shi_Heng: <p>Hi, This OOM exception comes from the python api implement of conv2d_weight actually.<br>
In backprop weight calculation, the output gradients need to be expanded with output channel times. When default cudnn implement this with data prefetch block and block (not allocate more memory), python api uses a repeat that will allocate a huge size of memory on output gradients tensor with unnecessary duplication of data.<br>
you can easily fix this by convert the repeat into a loop function at conv2d_weight.</p>

Mostafa_Elhoushi: <p>I think the expression for <code>grad_bias</code> should be fixed to:</p>
<pre><code class="lang-auto">grad_bias = grad_output.sum((0,2,3)).squeeze(0)
</code></pre>

Mostafa_Elhoushi: <p>Thanks for <a class="mention" href="/u/hanspinckaers">@hanspinckaers</a> for sharing this. I managed to run the code, however it is still slow. Much slower that using Pytorch’s <code>conv2d</code> and letting autograd do the work.</p>

hanspinckaers: <p>Hi Mostafa, in my benchmarks it does seem to perform at equal speed. how did you benchmark it?</p>

fsds: <p>If you mean running cudnn_convolution_backward_input() and cudnn_convolution_backward_weight() is slower than calling conv2d and letting autograd do the back-propagation, then it makes sense because you’re now calling two functions separately, and in addition (based on your previous comment) calculating the grad_bias.<br>
If you want to override the whole back-propagation process of Conv2d and still have the same processing time, you should use the combined cudnn_convolution_backward() that returns gradients w.r.t the input, gradients w.r.t the weights and gradients w.r.t the biases in that order.</p>
<p>The question here and <a class="mention" href="/u/hanspinckaers">@hanspinckaers</a> solution refer to overriding only torch.nn.grad.conv2d_weight, which is very expensive in memory, with cudnn_convolution_backward_weight().</p>

Mostafa_Elhoushi: <p>Thanks <a class="mention" href="/u/fsds">@fsds</a> and <a class="mention" href="/u/hanspinckaers">@hanspinckaers</a>!</p>
<p><a class="mention" href="/u/hanspinckaers">@hanspinckaers</a>: Thank you following up. it has been a while since I did the benchmark. I recall I was training ResNet18 on Imagenet. Using Pytorch’s <code>torch.nn.grad.conv2d_input(...)</code> and <code>torch.nn.grad.conv2d_weight(...)</code> was probably twice as slow and using twice as much memory than letting PyTorch derive the backward pass of Conv2d automatically.<br>
When I tried to use the method you provided in this <a href="https://discuss.pytorch.org/t/cuda-error-with-cudnn-convolution-backward-weight-function/41214/2">link</a>, made things a bit faster, but still much slower than PyTorch’s automatic backward pass.</p>
<p><a class="mention" href="/u/fsds">@fsds</a>:  Thanks for your answer. Are you referring to <code>td::tuple&lt;at::Tensor,at::Tensor&gt; cudnn_convolution_backward(...)</code> in:<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/Conv.cpp#L1064" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/Conv.cpp#L1064" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/Conv.cpp#L1064</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="1054" style="counter-reset: li-counter 1053 ;">
<li>    bool benchmark, bool deterministic)</li>
<li>{</li>
<li>  TensorArg grad_output{ grad_output_t, "grad_output", 1 },</li>
<li>            weight{ weight_t, "weight", 2 };</li>
<li>  return cudnn_convolution_backward_input(</li>
<li>      "cudnn_convolution_backward_input",</li>
<li>      input_size, grad_output, weight,</li>
<li>      padding, stride, dilation, groups, benchmark, deterministic);</li>
<li>}</li>
<li>
</li>
<li class="selected">std::tuple&lt;at::Tensor,at::Tensor&gt; cudnn_convolution_backward(</li>
<li>    const at::Tensor&amp; input, const at::Tensor&amp; grad_output_t, const at::Tensor&amp; weight,</li>
<li>    IntArrayRef padding, IntArrayRef stride, IntArrayRef dilation, int64_t groups,</li>
<li>    bool benchmark, bool deterministic, std::array&lt;bool,2&gt; output_mask) {</li>
<li>
</li>
<li>  Tensor grad_output = grad_output_t.contiguous(input.suggest_memory_format());</li>
<li>
</li>
<li>  Tensor grad_input, grad_weight;</li>
<li>  if (input.numel() == 0) {</li>
<li>    if (output_mask[0]) {</li>
<li>      grad_input = at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT);</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
So I just need to create a Python wrapper to it and invoke it in our backward pass?</p>

fsds: <p>Yes, that is what I’m referring to, although something is a little weird here, because in Pytorch1.2 (the version I’m currently using) this function returns 3 tensors (<code>std::tuple&lt;at::Tensor,at::Tensor,at::Tensor&gt;</code>, grad_output, grad_weight, grad_bias) and in the current state of the repository it returns only 2 (<code>std::tuple&lt;at::Tensor,at::Tensor&gt;</code>, grad_output and grad_weight).<br>
You can see here the implementation of 1.2:<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/v1.2.0/aten/src/ATen/native/cudnn/Conv.cpp#L39" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.2.0/aten/src/ATen/native/cudnn/Conv.cpp#L39" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/v1.2.0/aten/src/ATen/native/cudnn/Conv.cpp#L39</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="29" style="counter-reset: li-counter 28 ;">
<li>    IntArrayRef padding, IntArrayRef stride, IntArrayRef dilation, int64_t groups,</li>
<li>    bool benchmark, bool deterministic) {</li>
<li>  AT_ERROR("cudnn_convolution_backward_weight: ATen not compiled with cuDNN support");</li>
<li>}</li>
<li>
</li>
<li>at::Tensor cudnn_convolution_backward_bias(</li>
<li>    const at::Tensor&amp; grad_output) {</li>
<li>  AT_ERROR("cudnn_convolution_backward_bias: ATen not compiled with cuDNN support");</li>
<li>}</li>
<li>
</li>
<li class="selected">std::tuple&lt;at::Tensor,at::Tensor,at::Tensor&gt; cudnn_convolution_backward(</li>
<li>    const at::Tensor&amp; input, const at::Tensor&amp; grad_output, const at::Tensor&amp; weight,</li>
<li>    IntArrayRef padding, IntArrayRef stride, IntArrayRef dilation, int64_t groups,</li>
<li>    bool benchmark, bool deterministic, std::array&lt;bool,3&gt; output_mask) {</li>
<li>  AT_ERROR("cudnn_convolution_backward: ATen not compiled with cuDNN support");</li>
<li>}</li>
<li>
</li>
<li>at::Tensor cudnn_convolution_transpose(</li>
<li>    const at::Tensor&amp; input, const at::Tensor&amp; weight, const at::Tensor&amp; bias /* optional */,</li>
<li>    IntArrayRef padding, IntArrayRef output_padding, IntArrayRef stride, IntArrayRef dilation,</li>
<li>    int64_t groups, bool benchmark, bool deterministic) {</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Either way, using this should give you better performance.</p>

ptrblck: <p>We removed the <code>bias</code>, as the backward pass was faster with native PyTorch ops and had some other advantages as seen in <a href="https://github.com/pytorch/pytorch/pull/31524">this PR</a>.</p>

r3coder: <p>This doesn’t works if the stride and padding is different from basic values, so I’ve edited a bit.</p>
<blockquote>
<p>ValueError: requested an input grad size of [4, 4], but valid sizes range from [6, 6] to [6, 6] (for a grad_output of torch.Size([4, 4]))</p>
</blockquote>
<p>So, I saved some arguments with save_for_backward to work with it.</p>
<pre><code class="lang-auto">class Conv2dFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
        # Save arguments to context to use on backward
        # WARNING : if stride, padding, dilation etc is array, this will not work properly!!!!
        confs = torch.from_numpy(np.array([stride, padding, dilation, groups]))
        ctx.save_for_backward(input, weight, bias, confs)

        # Compute Convolution
        return F.conv2d(input, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
    
    @staticmethod
    def backward(ctx, grad_output):
        # Load saved tensors
        input, weight, bias, confs = ctx.saved_variables
        confs = confs.numpy()
        stride, padding, dilation, groups= confs[0], confs[1], confs[2], confs[3]

        # Calculate Gradient
        grad_input = grad_weight = grad_bias = None
        if ctx.needs_input_grad[0]:
            grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output, stride, padding, dilation, groups)
            
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output, stride, padding, dilation, groups)
                
        # WARNING : Bias maybe buggy, remove if it is buggy
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0).squeeze(0)


        # WARNING : Bias maybe buggy, remove if it is buggy
        if bias is not None:
            return grad_input, grad_weight, grad_bias, None, None, None, None
        else:
            return grad_input, grad_weight, None, None, None, None, None
</code></pre>
<p>Since there are more stride, padding, etc on the forward input, there need more output None to calculate grad.</p>
<p>Is works fine, but is there any elegant way to do this without returning useless None(s)?</p>
<p>Also, I think conv with bias is buggy. I’ll fix if I have more spare time</p>