Custom Autograd Function Backward pass not Called

eduardo4jesus: <p>So far, I’ve check all these posts that I found on the issue, but <code>None</code> of them seems to fit my case.</p>
<ul>
<li><a href="https://discuss.pytorch.org/t/backward-not-called-bugs/12380">backward-not-called-bugs</a></li>
<li><a href="https://discuss.pytorch.org/t/backward-not-called/9043/">backward-not-called</a></li>
<li><a href="https://discuss.pytorch.org/t/custom-autograd-function-backward-pass-not-called/9144">custom-autograd-function-backward-pass-not-called</a></li>
<li><a href="https://discuss.pytorch.org/t/using-backward-of-custom-autograd-function-with-loss-backward/67873">using-backward-of-custom-autograd-function-with-loss-backward</a></li>
</ul>
<p><strong>Here is my code</strong></p>
<pre><code class="lang-auto">class _Conv2dCustom(autograd.Function):

  # Kernel has to be of odd size
  @staticmethod
  def forward(ctx, input, weight, bias=None, padding=0):
    print('########## _Conv2dCustom.forward')
    input = input.clone() 
    ctx.save_for_backward(input, weight, bias)
    padding = (padding, padding) if type(padding) is int else padding
    # input = input.clone()
    output = _Conv2dCustom._conv2d(input, weight, padding)
    if bias is not None:
      output += bias.view(-1, 1, 1)
    return output
    
  @staticmethod
  def backward(ctx, grad_output):
    print('########## _Conv2dCustom.backward')
    # Kernel has to be of odd size
    input, weight, bias = ctx.saved_tensors
    grad_input = grad_weight = grad_bias = None
    if ctx.needs_input_grad[0]:
      # With padding, conv produce same size grad_output
      padding = weight.size(-1)//2
      # Difference to produce same size as input
      padding += (grad_output.size(-1)-input.size(-1))//2
      weight_t = torch.transpose(weight, -1, -2)
      grad_input = _Conv2dCustom._conv2d(grad_output, weight_t, padding)
    if ctx.needs_input_grad[1]:
        grad_weight = _Conv2dCustom._conv2d(input, grad_output, padding_used)
    if bias is not None and ctx.needs_input_grad[2]:
        grad_bias = grad_output.sum([0, 2, 3])

    return grad_input, grad_weight, grad_bias

  @staticmethod
  def _conv2d(input, weight, padding=(0, 0), stride=(1, 1)):
    output_h = (input.size(-2) - weight.size(-2) + 2*padding[0])//stride[0] + 1
    output_w = (input.size(-1) - weight.size(-1) + 2*padding[1])//stride[1] + 1

    Xunfold = F.unfold(input, weight.size()[2:], padding=padding, stride=stride)
    weight_flat = weight.view(weight.size(0), -1)
    conv = weight_flat@Xunfold
    output = conv.view(input.size(0), weight.size(0), output_h, output_w)
    return output

class Conv2dCustom(nn.Conv2d):
  # Torch master as in Feb 23rd 2020.
  def _conv_forward(self, *args, **kwargs):
    raise NotImplementedError('This Method Is Not Implemented')

  # Torch v1.4.0  
  def conv2d_forward(self, input, weight):
    conv2d = _Conv2dCustom.apply
    conv = conv2d(input, weight, self.bias, self.padding)
    return conv
</code></pre>
<p><strong>Output</strong></p>
<pre><code class="lang-auto">==========
Started Training
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
==========
Finished Training
</code></pre>
<p>In my Network, I have two of these <code>Conv2dCustom</code> layers, and I also put hooks on each layer. I noticed that the backward hook for all my layers, except the most inner <code>Conv2dCustom</code> layer.</p>

albanD: <p>Can you share <code>_Conv2dCustom._conv2d()</code> ?</p>
<p>Also can you write a small code sample that uses the <code>Conv2dCustom</code> that triggers this so that we can reproduce locally?</p>

eduardo4jesus: <p><code>_Conv2dCustom._conv2d()</code>  is in the code above, in the question. But here it is.</p>
<pre><code class="lang-auto">  @staticmethod
  def _conv2d(input, weight, padding=(0, 0), stride=(1, 1)):
    output_h = (input.size(-2) - weight.size(-2) + 2*padding[0])//stride[0] + 1
    output_w = (input.size(-1) - weight.size(-1) + 2*padding[1])//stride[1] + 1

    Xunfold = F.unfold(input, weight.size()[2:], padding=padding, stride=stride)
    weight_flat = weight.view(weight.size(0), -1)
    conv = weight_flat@Xunfold
    output = conv.view(input.size(0), weight.size(0), output_h, output_w)
    return output
</code></pre>
<hr>
<h2>Reproducing with a small test</h2>
<p>Code</p>
<pre><code class="lang-auto">image = np.random.randint(255, size=(1, 1, 32, 32)).astype('float32')
image = torch.from_numpy(image).requires_grad_(True)
layer = Conv2dCustom(1, 1, 5)
result = layer.forward(image)
</code></pre>
<p>or</p>
<pre><code class="lang-auto">image = np.random.randint(255, size=(1, 1, 32, 32)).astype('float32')
image = torch.autograd.Variable(torch.from_numpy(image)).requires_grad_(True)
layer = Conv2dCustom(1, 1, 5)
result = layer.forward(image)
</code></pre>
<p>Output</p>
<pre><code class="lang-auto">########## _Conv2dCustom.forward
</code></pre>
<hr>
<p>Code</p>
<pre><code class="lang-auto">dout= np.random.randint(255, size=(1, 1, 28, 28)).astype('float32')
dout = torch.from_numpy(dout)
result.backward(dout)
</code></pre>
<p>Nothing is printed.</p>

albanD: <p>Sorry I missed your conv2d as a static function, my bad.</p>
<p>Few details:</p>
<ul>
<li>You nn.Module does not work for me and through the not implemented error from <code>_conv_forward</code>
</li>
<li>You should never call the <code>.forward()</code> of the nn.Module directly. You should use the <code>__call__</code> on it as  <code>result = layer(image)</code>.</li>
<li>You don’t need to wrap things in Variables anymore. You can just use regular Tensors: <code>image = torch.from_numpy(image).requires_grad_()</code>.</li>
</ul>
<p>By removing the custom module, it runs.<br>
This is a know issue and if you use the nightly build, you will see the following warning:<br>
“”"<br>
…/torch/csrc/autograd/variable.cpp:421: UserWarning: Output 0 of _Conv2dCustomBackward is a view and its base or another view of its base has been modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is deprecated and will be forbidden starting version 1.6. You can remove this warning by cloning the output of the custom Function.<br>
“”"</p>
<p>As it states, the fact that your custom Function returns a view and that you modify it inplace in when adding the bias break some internal autograd assumptions.<br>
You should either change <code>_conv2d</code> to <code>return output.clone()</code> to avoid returning a view. Or change your bias update to <code>output = output + bias.view(-1, 1, 1)</code> to avoid the inplace operations.</p>

eduardo4jesus: <p>No worries, that happens to me a lot</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71065">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Sorry I missed your conv2d as a static function, my bad.</p>
</blockquote>
</aside>
<hr>
<p>I am working on Colab, which uses the v1.4.0. I put that exception to get the error in case the update the current stable version.</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71065">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>You nn.Module does not work for me and through the not implemented error from <code>_conv_forward</code></p>
</blockquote>
</aside>
<hr>
<p>Oh I see. Thanks. On my original code, the one I first experimented this issue, I don’t call <code>forward()</code> myself. I use the class from standard torch examples of a neural network (Lenet) with a few changes.</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71065">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>You should never call the <code>.forward()</code> of the nn.Module directly. You should use the <code>__call__</code> on it as <code>result = layer(image)</code> .</p>
</blockquote>
</aside>
<hr>
<p>I suspected. Thanks for the clarification. I did both cases just to be safe, as a quick test.</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71065">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>You don’t need to wrap things in Variables anymore. You can just use regular Tensors: <code>image = torch.from_numpy(image).requires_grad_()</code> .</p>
</blockquote>
</aside>
<hr>
<p>Wow. I really looked for inplace operations. At first, I thought <code>output += bias.view()</code><br>
could be inplace. But, for some reason, I sticked with “Well, that should be the same thing as <code>output = output+bias</code>”. I totally should had tried that!</p>
<aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71065">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>“”"<br>
…/torch/csrc/autograd/variable.cpp:421: UserWarning: Output 0 of _Conv2dCustomBackward is a view and its base or another view of its base has been modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is deprecated and will be forbidden starting version 1.6. You can remove this warning by cloning the output of the custom Function.<br>
“”"</p>
<p>As it states, the fact that your custom Function returns a view and that you modify it inplace in when adding the bias break some internal autograd assumptions.<br>
You should either change <code>_conv2d</code> to <code>return output.clone()</code> to avoid returning a view. Or change your bias update to <code>output = output + bias.view(-1, 1, 1)</code> to avoid the inplace operations.</p>
</blockquote>
</aside>
<hr>
<p>Thank you so much! <strong>You saved my week!</strong> I really appreciate this community!</p>
<hr>
<p>Just one more question, from the <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">Extending Pytorch page</a>:</p>
<h4><a name="shouldnt-the-inplace-operation-be-an-issue-only-for-the-input-parameters-on-the-forward-pass-1" class="anchor" href="#shouldnt-the-inplace-operation-be-an-issue-only-for-the-input-parameters-on-the-forward-pass-1"></a>Shouldn’t the <code>inplace</code> operation be an issue only for the input parameters on the forward pass?</h4>
<blockquote>
<ul>
<li><a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.function._ContextMethodMixin.mark_dirty" rel="noopener nofollow ugc"> <code>mark_dirty()</code> </a> must be used to mark any input that is modified inplace by the forward function.</li>
</ul>
</blockquote>

albanD: <aside class="quote no-group" data-username="eduardo4jesus" data-post="5" data-topic="71065">
<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/eduardo4jesus/48/20203_2.png" class="avatar"> eduardo4jesus:</div>
<blockquote>
<p>Shouldn’t the <code>inplace</code> operation be an issue only for the input parameters on the forward pass?</p>
</blockquote>
</aside>
<p>This is a different thing. If an input is modified inplace, we need to handle it specifically and so we need the user to tell us they did it.</p>
<p>But that logic that handle inplace needs to look at how the view was created. And unfortunately, can only handle a subset of these (they are fairly corner case). We will now raise a nice error if you do so.</p>