torch.autograd.Function overwrite

username: Dale_Song
cooked: <p>hi, when inherit <code>autograd.Function</code> to define an op e.g.</p>
<pre><code class="lang-auto">class MyFunc(autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        some_inplace_operation1(x)
        return x
   @staticmethod
    @once_differentiable
    def backward(ctx, dy):
        some_inplace_operation2(dy)
        return dy
</code></pre>
<p>in some cases, there will be total error gradient (depend on the followed op). What’s the reason? A concrete example is <a href="https://github.com/mapillary/inplace_abn/issues/155" rel="nofollow noopener">here</a>.</p>

username: albanD
cooked: <p>Hi,</p>
<p>You are not allowed to modify <code>dy</code> inplace in the backward function. If you do so, the computed gradients can be arbitrarily wrong <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Dale_Song
cooked: <p>Could you please explain more why the computed gradients can be arbitrarily wrong and is  there a solution to safely modify  <code>dy</code> because this can save memory and improve efficiency.</p>

username: albanD
cooked: <p>Because this Tensor is supposed to contain the value of the gradient for a given Tensor in the forward pass.<br>
But this value might be needed in other places. If you modify it inplace, these other places will see wrong values.</p>

username: Dale_Song
cooked: <p>Where may it be used, assume that op following this op has it as the only one input.</p>

username: albanD
cooked: <p>In general, it should never <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"> In some limited cases, it might work, but this is not guaranteed.<br>
Also your forward function is missing a call to <code>mark_dirty()</code>. See the <a href="https://pytorch.org/docs/master/notes/extending.html">doc</a> about extending the autograd.</p>

username: Dale_Song
cooked: <p>I did call <code>mark_dirty</code> <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"> Is there a way to find out what kind of cases it should be?</p>

username: albanD
cooked: <p>Could you give a full code sample that reproduce the error you’re seeing?</p>

username: Dale_Song
cooked: <p>You can try the <a href="https://gist.github.com/knsong/daea712fa06043ccb501ad1c92143bfa" rel="noopener nofollow ugc">resnet18</a> with the op <code>InPlaceABN</code> at this <a href="https://github.com/mapillary/inplace_abn/tree/826d7c99d7c6494b95e1c16cda6f37afcd8c7190" rel="noopener nofollow ugc">commit</a> . I train it on cifar10 using code <a href="https://github.com/kuangliu/pytorch-cifar" rel="noopener nofollow ugc">here</a> and pytorch 1.1.0, its validation accuracy is always lower than that of resnet18 with standard <code>BatchNorm2d</code> and there is a significant accuray drop during training:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/b/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1" title="result"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/b/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1_2_690x379.jpeg" alt="result" data-base62-sha1="mbHF0UItinak81KdEvfVx6lbQcx" width="690" height="379" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/b/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1_2_690x379.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/b/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1_2_1035x568.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/b/9b828701f8cb22d09366dfa6ae34fda0bf5c95d1_2_1380x758.jpeg 2x" data-dominant-color="292828"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">result</span><span class="informations">2048×1126 369 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
How to fix this problem without modify the op since it is more memory efficient?</p>

username: albanD
cooked: <p>That’s a lot of code… that would really help if you had a small code sample that we can run and see what it is doing easily.</p>
<p>Are you sure the accuracy drop comes from the autograd? Can you try temporarily adding a <code>dy.clone()</code> at the beginning of your function and see if you still have the same behavior?<br>
Such behavior in the loss can also come from instability of your new operation.</p>

username: Dale_Song
cooked: <p>Sorry for the plenty of code, Let me make a minimal one later. I have tried that(adding a  <code>dy.clone()</code>  at the beginning of backward) and there was no accuracy drop and a higher final accuracy(5 percent point higher).</p>

username: Dale_Song
cooked: <p>Hi, here is a <a href="https://gist.github.com/knsong/1b2aeaf1cdfef28d52138df5d05cd949" rel="nofollow noopener">minimal sample code</a> that can reproduce the error. Please check it.<br>
I find out that the error is related to at <code>out = out + self.shortcut(x)</code> when  <code>len(self.shortcut) == 0</code> and <code>out</code> is outputed by the inplace op,  what’s the reason behind it?</p>
<pre><code class="lang-auto">class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*planes)
            )
            self.bn2 = nn.BatchNorm2d(planes)
        else:
            #self.n2 = nn.BatchNorm2d(planes) #use name `n2` to keep it `nn.BatchNorm2d` when `len(self.shortcut) == 0`
                                             #can avoid the gradient error, but why?
            self.bn2 = nn.BatchNorm2d(planes) #in this case, 'bn2' will be repaced with `my_func*` when `len(self.shortcut) == 0`
                                             # and will cause the gradient error.


    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.conv2(out)
        if len(self.shortcut):
            out = self.bn2(out)
        else:
            out = self.bn2(out)
            #out = self.n2(out)
        out = out + self.shortcut(x)
        out = F.relu(out, inplace=True)
        return out
</code></pre>

username: albanD
cooked: <p>There is no custom autograd Function in this code right?<br>
This code looks fine, what is the exact issue you’re seeing with it?</p>

username: Dale_Song
cooked: <p>The full same code is: <a href="https://gist.github.com/knsong/1b2aeaf1cdfef28d52138df5d05cd949" rel="nofollow noopener">https://gist.github.com/knsong/1b2aeaf1cdfef28d52138df5d05cd949</a> , and my custom autograd Function is there, and if you execute it with  <code>python backward_inplace.py</code> , you will see the gradient error between the two version(inplace or no inplace) of the Function, like <code>grad x error max: tensor(0.8613)</code>.</p>

username: albanD
cooked: <p>As I said, you are not allowed to modify the grad inplace. That will make what the hooks (and potentially other function) see. So this is expected that you get wrong gradients if you modify the grad inplace in the backward function !</p>

username: Dale_Song
cooked: <p>It does’t always result in wrong gradients as I expalined in the sample code, only at <a href="https://gist.github.com/knsong/1b2aeaf1cdfef28d52138df5d05cd949#file-backward_inplace-py-L97" rel="nofollow noopener">here</a> it will do. So I will greatly appreciate if there is some details about what happened here?</p>

username: albanD
cooked: <p>We would need to dig into how each of these functions are implemented.<br>
In particular we need to find out what uses this particular gradient buffer and in which order.</p>
<p>My point was more that it is a dangerous thing to do. Even if you find why it works in this case, for this version of pytorch, your code might silently be broken by the next version of pytorch if we change some internal implementations of Functions or the autograd engine <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>