How to fix this nan bug?

username: JWarlock
cooked: <p>Hi, I’m trying to modify the mean/std of one feature with the mean/std calculated from another feature. It looks like this (certain simplification is made since original code is much more complicated)</p>
<pre><code class="lang-auto">
def exchange(vs, vt):
    # vs and vt are of the same size NxCxHxW
    vs_mean = torch.mean(vs, dim=(2, 3))
    vs_std = torch.std(vs, dim=(2, 3)) + self.eps
    raw_es_domain = torch.cat((vs_mean, vs_std), dim=1)
    es_domain = self.relu(self.fc1(raw_es_domain))
    es_domain = self.relu(self.fc2(es_domain))

    weight = self.fc_w(es_domain) # size of NxCx1x1
    bias = self.fc_b(es_domain) # size of NxCx1x1

    vt = nn.InstanceNorm2d(vt.size(1), affine=False)(vt)
    vt = weight * vt + bias

    return vt
</code></pre>
<p>However, there is this nan bug that happens sometimes, while sometimes this code snippet works just fine. I’ve used <code>torch.autograd.detect_anomaly()</code> to debug, and it told me the <code> torch.std</code> operation was the culprit. The simplified log information is listed as follows:</p>
<pre><code class="lang-auto">/opt/conda/conda-bld/pytorch_1573049306803/work/torch/csrc/autograd/python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
vs_std = torch.std(vs, dim=(2, 3)) + self.eps
loss.backward()
File "/home/user/anaconda3/envs/pytorch-1.3.1/lib/python3.7/site-packages/torch/tensor.py", line 166, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/user/anaconda3/envs/pytorch-1.3.1/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
</code></pre>
<p>But I’ve already added an <code>eps</code> to the <code>vs_std</code>. Is there any way to fix this bug? Or is there any other way to find out what’s the cause of nan? Thanks!</p>

username: tom
cooked: <p>This is likely because vs is constant:</p>
<pre><code class="lang-python">a = torch.zeros(10, requires_grad=True)                                                                                                                                                                    
a.std().backward()                                                                                                                                                                                         
a.grad                                                                                                                                                                                                     
</code></pre>
<p>gives<br>
<code>tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])</code></p>
<p>What you could do is add a gradient hook and modify the gradient to replace the nan with 0 (using torch.where).</p>
<p>Best regards</p>
<p>Thomas</p>

username: JWarlock
cooked: <p>Thanks! Thomas</p>
<p>That is indeed the case for me. I checked the <code>vs_std</code> tensor and found something like <code>vs[14, 205, :, :]</code> was an all-zero tensor. So I assmue modify the nan gradient of <code>vs</code> with 0 would solve the problem.</p>
<p>However, I’ve tried to add gradient hook to both vs and vt (or vs_std and vt_std), and it didn’t work.</p>
<pre><code class="lang-auto">def hook_fn_backward(grad):
    grad = torch.where(torch.logical_not(torch.isnan(grad)), grad, torch.zeros_like(grad).to(grad.device))
return grad
</code></pre>
<pre><code class="lang-auto">vs.register_hook(hook_fn_backward)
vs_mean.register_hook(hook_fn_backward)
vs_std.register_hook(hook_fn_backward)
</code></pre>
<p>It still gives following error<br>
<code>RuntimeError: Function 'StdBackward1' returned nan values in its 0th output.</code></p>
<p>The strange thing is, I tried to set trace in <code>hook_fn_backward</code>, but I only found that <code>torch.any(torch.isnan(grad)) == False</code>.  So where should I add this hook exactly?</p>

username: JWarlock
cooked: <p>After Further debugging, I find that add a gradient hook to <code>vs</code> and modify the gradient to replace the nan with 0 <strong>does</strong>  solve the problem mentioned above. That is to say, the <code>nan</code> gradient from<code> torch.std()</code> is replaced with 0.</p>
<p>However, I then found there is another nan bug in this code. And since I’m using <code>torch.autograd.detect_anomaly()</code> to find out which line is the culprit, the line<br>
<code>vs_std = torch.std(vs, dim=(2, 3)) + self.eps</code><br>
always shows up in the error log. It seems that <code>torch.autograd.detect_anomaly()</code> still gives an error at this line even after adding a gradient hook to <code>vs</code> (in fact, I have to comment out <code>torch.autograd.detect_anomaly()</code> to make the code run normally). And that covers up the trace of the second nan bug.</p>
<p>What a code snippet! I never thought of another nan bug, so naturally I’ve wasted lots of time.</p>
<p>Here is the thing: I’ve added gradient hooks to every module in this custom block, and I find the <code>linear</code> layers are also having nan gradient. There are many <code>linear</code> layers receiving nan <code>grad_out</code> and sending <code>grad_in</code>, but there is also this one ‘linear’ layer ‘fc_w’ <strong>receving normal ‘grad_out’ and sending nan ‘grad_in’</strong>.  At first, I thought it was the input <code>es_domain</code> was nan. But then I found out <code>es_domain</code> became nan at iteration 1, while nan gradient in <code>linear</code> layer happened at iteration 0.</p>
<p>So I wonder how to fix this bug? Or is there any way to inpect the internal backprop procedure in a <code>linear</code> layer?</p>
<p>Thanks again!</p>

username: JWarlock
cooked: <p><a class="mention" href="/u/tom">@tom</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/alband">@albanD</a>  Hi, sorry for the tagging thing, but I didn’t receive any response, so i have to tag you.</p>
<p>Like I said, after addin a gradient hook to<code> vs</code>, the gradient in backbone (all those conv-bn-relu layers) are now normal.  But I find the gradient of linear layers in this custom block contains nan (this block is inserted in the middle stage of backbone, like after stage 3 of ResNet or relu3-1 of VGGNet).  After optimization step, the weight of these linear layers become nan. And in the next iteration, the model output become nan too.</p>
<pre><code class="lang-auto"> for name, param in self.model.backbone.named_parameters():
    print(name, torch.isnan(param.grad).any())
###### output
conv1.weight tensor(False, device='cuda:0')
bn1.weight tensor(False, device='cuda:0')
bn1.bias tensor(False, device='cuda:0')
layer1.0.conv1.weight tensor(False, device='cuda:0')
layer1.0.bn1.weight tensor(False, device='cuda:0')
layer1.0.bn1.bias tensor(False, device='cuda:0')
layer1.0.conv2.weight tensor(False, device='cuda:0')
layer1.0.bn2.weight tensor(False, device='cuda:0')
layer1.0.bn2.bias tensor(False, device='cuda:0')
layer1.1.conv1.weight tensor(False, device='cuda:0')
layer1.1.bn1.weight tensor(False, device='cuda:0')
layer1.1.bn1.bias tensor(False, device='cuda:0')
layer1.1.conv2.weight tensor(False, device='cuda:0')
layer1.1.bn2.weight tensor(False, device='cuda:0')
layer1.1.bn2.bias tensor(False, device='cuda:0')
layer2.0.conv1.weight tensor(False, device='cuda:0')
layer2.0.bn1.weight tensor(False, device='cuda:0')
layer2.0.bn1.bias tensor(False, device='cuda:0')
layer2.0.conv2.weight tensor(False, device='cuda:0')
layer2.0.bn2.weight tensor(False, device='cuda:0')
layer2.0.bn2.bias tensor(False, device='cuda:0')
layer2.0.downsample.0.weight tensor(False, device='cuda:0')
layer2.0.downsample.1.weight tensor(False, device='cuda:0')
layer2.0.downsample.1.bias tensor(False, device='cuda:0')
layer2.1.conv1.weight tensor(False, device='cuda:0')
layer2.1.bn1.weight tensor(False, device='cuda:0')
layer2.1.bn1.bias tensor(False, device='cuda:0')
layer2.1.conv2.weight tensor(False, device='cuda:0')
layer2.1.bn2.weight tensor(False, device='cuda:0')
layer2.1.bn2.bias tensor(False, device='cuda:0')
layer3.0.conv1.weight tensor(False, device='cuda:0')
layer3.0.bn1.weight tensor(False, device='cuda:0')
layer3.0.bn1.bias tensor(False, device='cuda:0')
layer3.0.conv2.weight tensor(False, device='cuda:0')
layer3.0.bn2.weight tensor(False, device='cuda:0')
layer3.0.bn2.bias tensor(False, device='cuda:0')
layer3.0.downsample.0.weight tensor(False, device='cuda:0')
layer3.0.downsample.1.weight tensor(False, device='cuda:0')
layer3.0.downsample.1.bias tensor(False, device='cuda:0')
layer3.1.conv1.weight tensor(False, device='cuda:0')
layer3.1.bn1.weight tensor(False, device='cuda:0')
layer3.1.bn1.bias tensor(False, device='cuda:0')
layer3.1.conv2.weight tensor(False, device='cuda:0')
layer3.1.bn2.weight tensor(False, device='cuda:0')
layer3.1.bn2.bias tensor(False, device='cuda:0')
layer4.0.conv1.weight tensor(False, device='cuda:0')
layer4.0.bn1.weight tensor(False, device='cuda:0')
layer4.0.bn1.bias tensor(False, device='cuda:0')
layer4.0.conv2.weight tensor(False, device='cuda:0')
layer4.0.bn2.weight tensor(False, device='cuda:0')
layer4.0.bn2.bias tensor(False, device='cuda:0')
layer4.0.downsample.0.weight tensor(False, device='cuda:0')
layer4.0.downsample.1.weight tensor(False, device='cuda:0')
layer4.0.downsample.1.bias tensor(False, device='cuda:0')
layer4.1.conv1.weight tensor(False, device='cuda:0')
layer4.1.bn1.weight tensor(False, device='cuda:0')
layer4.1.bn1.bias tensor(False, device='cuda:0')
layer4.1.conv2.weight tensor(False, device='cuda:0')
layer4.1.bn2.weight tensor(False, device='cuda:0')
layer4.1.bn2.bias tensor(False, device='cuda:0')
layer3_custom_block0.fc1.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc2.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_w.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_w.bias tensor(False, device='cuda:0')
layer3_custom_block0.fc_b.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_b.bias tensor(False, device='cuda:0')
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(True, device='cuda:0')

</code></pre>
<p>P.S.  The mean/std of new vt is also calculated, which is omitted in the first post. I’ve also added gradient hook to <code>new_vt</code>, but the nan gradient in fc layer still exist.</p>
<pre><code class="lang-auto">def exchange(vs, vt):
    # vs and vt are of the same size NxCxHxW
    vs_mean = torch.mean(vs, dim=(2, 3))
    vs_std = torch.std(vs, dim=(2, 3)) + self.eps
    raw_es_domain = torch.cat((vs_mean, vs_std), dim=1)
    es_domain = self.relu(self.fc1(raw_es_domain))
    es_domain = self.relu(self.fc2(es_domain))

    weight = self.fc_w(es_domain) # size of NxCx1x1
    bias = self.fc_b(es_domain) # size of NxCx1x1

    vt = nn.InstanceNorm2d(vt.size(1), affine=False)(vt)
    new_vt = weight * vt + bias

    new_vt_mean = torch.mean(new_vt, dim=(2, 3))
    new_vt_std = torch.std(new_vt, dim=(2, 3)) + self.eps
    raw_new_et_domain = torch.cat((new_vt_mean, new_vt_std), dim=1)
    new_et_domain = self.relu(self.fc1(raw_new_et_domain))
    new_et_domain = self.relu(self.fc2(new_et_domain))

    return vt, es_domain, new_et_domain
</code></pre>
<p>The funny thing is,  after I add hook to moniter the gradient of <code>weight</code> and <code>bias</code> in <code>custom_block1</code>, the nan in fc_w of custom_block1 just disappear, but the nan in fc_b do not!</p>
<pre><code class="lang-auto">def hook(grad):
    if torch.any(torch.isnan(grad)):
        print('fixing nan gradient')
    grad = torch.where(torch.logical_not(torch.isnan(grad)), grad, torch.zeros_like(grad).to(grad.device))
    return grad

weight.register_hook(grad)
bias.register_hook(grad)

####### after above hook registered, check the gradient of custom_block 
layer3_custom_block0.fc1.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc2.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_w.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_w.bias tensor(False, device='cuda:0')
layer3_custom_block0.fc_b.weight tensor(False, device='cuda:0')
layer3_custom_block0.fc_b.bias tensor(False, device='cuda:0')
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(False, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(False, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(True, device='cuda:0')

</code></pre>
<p>I know it maybe a dummy question to ask, but how should I debug?</p>

username: tom
cooked: <p>Well, I don’t know if the hook keeps the anomaly mode from complaining and while I don’t fully understand what you are doing, but it would seem that std still is the critical bit.<br>
One (slow, but maybe OK for debugging) alternative to the gradient hook is to insert some noise <code>new_vt = new_vt + 1e-6 * torch.randn_like(new_vt)</code> in the forward.</p>
<p>Best regards</p>
<p>Thomas</p>

username: JWarlock
cooked: <p>Yes, it works!<br>
After inserting some noise to both <code>vs</code> and <code>new_vt</code> as you suggested (without gradient hook), this code can now runs normally without giving nan values.</p>
<p>What I’m trying to do is to modify the moment from features of target domain (vt) with the moment from features of source domain (vs), and calculate some loss based on modified feature <code>new_vt</code> and its moment  <code>new_et_domain</code>. In practice, there are multiple domains and each domain has its corresponding custom block. That is to say, <code>custom_block{i}.exchange(v_{i}, v_{j})</code> is performed on each pair of domains (for example, 3 domains in total).</p>
<p>However, there is one strange thing I don’t understand. If inserting noise works, then the nan should come from the std part. But why doesn’t gradient hook work as well? I tried to only add noise to <code>vs</code>, and it gives this error instead:</p>
<pre><code class="lang-auto">/opt/conda/conda-bld/pytorch_1573049306803/work/torch/csrc/autograd/python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
 new_vt = weight * vt + bias

RuntimeError: Function 'MulBackward0' returned nan values in its 0th output.
</code></pre>
<p>Also, if I add noise to <code>vs</code> and gradient hook to <code>new_vt</code>, the bias of <code>fc_b</code> becomes normal</p>
<pre><code class="lang-auto"> for name, param in self.model.backbone.named_parameters():
    print(name, torch.isnan(param.grad).any())

###### only noise to vs
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(True, device='cuda:0')

###### both noise to vs and gradient hook to new_vt
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(False, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(False, device='cuda:0')
</code></pre>
<p>So the nan gradient of <code>new_vt</code> comes from multiplication instead of std? ? The only reason I can see for multiplication to give nan gradient is the input is nan as well. But that doesn’t make sense since inserting noise to <code>new_vt</code> solves the problem. Is there another possible reason for multiplication to give nan gradient?</p>
<p>(P.S. I find this post <a href="https://github.com/pytorch/pytorch/issues/15288" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/15288</a>, but I’m not sure if it’s the same case)</p>

username: tom
cooked: <p>So if you run anomaly mode, the error might be thrown between the producing of NaN and before the hook is applied. (I must admit I don’t know and in lieu of a simple snippet I can copy-paste into PyTorch, I’m not going to start experiments, sorry.)<br>
One way around this might be to define your own <code>autograd.Function</code> for <code>std</code> instead of using the built-in one to not return NaNs in the first place.</p>

username: JWarlock
cooked: <p>Yes, I agree.</p>
<p>I must say returning nan gradient while computing std on all constant tensor is just so strange! The natural way is to return all 0 gradient, isn’t it?  Maybe PyTorch should fix it.</p>
<p>For others who have encountered the same problem, I write my own <code>CustomStd</code> like this:</p>
<pre><code class="lang-auto">class CustomStd(torch.autograd.Function):
@staticmethod
    def forward(ctx, input, dim, eps=1e-5, unbiased=True, keepdim=False):
        dev = input - input.mean(dim=dim, keepdim=True)
        ctx.save_for_backward(input)
        ctx.eps=eps
        ctx.dev = dev
        ctx.numdim = input.dim()
        ctx.numel = functools.reduce(lambda x, y: x * y, [input.size(d) for d in dim])
        if unbiased:
            ctx.numel -= 1
        ctx.std = torch.sqrt(torch.sum(dev * dev, dim=dim, keepdim=True) / ctx.numel)
        return ctx.std if keepdim else ctx.std.squeeze()

    @staticmethod
    def backward(ctx, grad_output):
        input,= ctx.saved_tensors
        grad_input = grad_output
        for i in range(grad_output.dim(), ctx.numdim):
            grad_input = grad_input.unsqueeze(i)
        grad_input = ctx.dev * (ctx.numel - 1) / (ctx.numel**2) / (ctx.std + ctx.eps) * grad_input
        return grad_input, None, None, None, None
</code></pre>
<p>It can pass a test like this:</p>
<pre><code class="lang-auto">input = torch.randn(16, 256, 14, 14, dtype=torch.float, requires_grad=True)
std1 = input.std(dim=(2,3))
std2 = CustomStd.apply(input, (2,3))
torch.sum(std1).backward(retain_graph=True)
grad1= input.grad
input.grad=None
torch.sum(std2).backward()
grad2 = input.grad
############### TEST
torch.allclose(std1, std2)       # True
torch.allclose(grad1, grad2, atol=1e-3)      # True
</code></pre>
<p>In the second test (gradient), I must set atol=1e-3 to make the test return true. So I guess there is a subtle difference between <code>CustomStd</code> and <code>torch.std</code>.<br>
Also, now  <code>CustomStd</code> can return 0 gradient for all constant tensor.</p>

username: tom
cooked: <aside class="quote no-group" data-username="JWarlock" data-post="10" data-topic="90291">
<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/jwarlock/48/20044_2.png" class="avatar"> JWarlock:</div>
<blockquote>
<p>I must say returning nan gradient while computing std on all constant tensor is just so strange! The natural way is to return all 0 gradient, isn’t it? Maybe PyTorch should fix it.</p>
</blockquote>
</aside>
<p>I don’t think it’s natural to want the gradient of std at constant vectors, but yes, you could argue that picking 0 as a subgradient is better than NaN (which is likely artificially there from writing it as var(…)**0.5).</p>