Gradient w.r.t. Inputs using BatchNorm

mpr: <p>Hello,</p>
<p>I’m trying to calculate the gradient of the output of a simple neural network with respect to the inputs. The result looks fine when I don’t use a BatchNorm layer. Once I do use it, the result doesn’t seem to make much sense. Below is a short example to reproduce the effect.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import matplotlib.pyplot as plt
</code></pre>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, batch_norm):
        super().__init__()
        
        self.batch_norm = batch_norm
        self.act_fn = nn.Tanh()
        
        self.aff1 = nn.Linear(1, 10)
        self.aff2 = nn.Linear(10, 1)
        
        if batch_norm:
            self.bn = nn.BatchNorm1d(10, affine=False)  # False for simplicity
        
    def forward(self, x):
        x = self.aff1(x)
        x = self.act_fn(x)
        
        if self.batch_norm:
            x = self.bn(x)
            
        x = self.aff2(x)
        return x
</code></pre>
<pre><code class="lang-auto">x_vals = torch.linspace(0, 1, 100)
x_vals.requires_grad = True

fig, axs = plt.subplots(ncols=2, figsize=(16, 5))

for seed, bn, ax1 in zip([11, 7], [False, True], axs):  # different seeds for better illustration of effect
    torch.manual_seed(seed)
    net = Net(batch_norm=bn)

    net.train()
    pred = net(x_vals[:, None])
    pred_dx = torch.autograd.grad(pred.sum(), x_vals, create_graph=True)[0]

    # visualization
    ax2 = ax1.twinx()

    ax1.plot(x_vals.detach(), pred.detach())
    ax2.plot(x_vals.detach(), pred_dx.detach(), linestyle='--', color='orange')

    min_idx = torch.argmin((pred[1:]-pred[:-1])**2)
    ax2.axvline(x_vals[min_idx].detach(), color='gray', linestyle='dotted')
    ax2.axhline(0, color='gray', linestyle='dotted')
    ax1.set_title(('With' if bn else 'Without') + ' Batch Norm')
    
plt.show()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/3/038f35df07bc3352c48c0837d1f8116b81fbc94c.png" data-download-href="https://discuss.pytorch.org/uploads/default/038f35df07bc3352c48c0837d1f8116b81fbc94c" title="bn_question_1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/038f35df07bc3352c48c0837d1f8116b81fbc94c_2_690x375.png" alt="bn_question_1" data-base62-sha1="vug4WQDMJNzUQJXGkGpzFLMyhm" width="690" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/038f35df07bc3352c48c0837d1f8116b81fbc94c_2_690x375.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/038f35df07bc3352c48c0837d1f8116b81fbc94c_2_1035x562.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/038f35df07bc3352c48c0837d1f8116b81fbc94c_2_1380x750.png 2x" data-dominant-color="FDFDFD"><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">bn_question_1</span><span class="informations">1920×1044 108 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></p>
<p>The result also seems to be fine when I use evaluation mode. Unfortunately I can’t just switch to eval() mode because the nature of my problem (PINNs) requires calculating these gradients during training.</p>
<p>This question is probably related to the post <a href="https://discuss.pytorch.org/t/the-gradients-of-batchnorm-layer-at-mode-of-model-train-and-model-eval/117000" class="inline-onebox">The gradients of BatchNorm layer at mode of model.train() and model.eval()</a> to which there are no answers.</p>
<p>I’m using python version 3.9.5, pytorch version 1.9.0+cu102.</p>
<p>Thanks for your help!</p>

mpr: <p>I think BatchNorm1d might not .detach() the mean and variance when normalizing. When I use the implementation below it seems work (simplified by removing the extra bias and weight).</p>
<pre><code class="lang-auto">class BatchNorm(nn.Module):
    def __init__(self, nFeatures, eps=1e-5, momentum=0.1):
        super().__init__()
        self.register_buffer("moving_avg", torch.zeros(nFeatures))
        self.register_buffer("moving_var", torch.ones(nFeatures))
        self.register_buffer("eps", torch.tensor(eps))
        self.register_buffer("momentum", torch.tensor(momentum))
    
    def forward(self, x):        
        if self.training:
            mean = x.mean(dim=0)
            var = x.var(dim=0)
            self.moving_avg = self.moving_avg * self.momentum + mean * (1 - self.momentum)
            self.moving_var = self.moving_var * self.momentum + var * (1 - self.momentum)

        else:
            mean = self.moving_avg
            var = self.moving_var
            
        # don't think original implementation uses .detach() on mean &amp; var
        x_norm = (x - mean.detach()) / (torch.sqrt(var.detach() + self.eps))
        return x_norm
</code></pre>
<p>I tried to find the exact implementation used by PyTorch to check. In <a href="https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#batch_norm" rel="noopener nofollow ugc">torch.nn.functional</a> the batch_norm function calls torch.batch_norm. Could somebody point me into the direction of where that’s implemented?</p>

vainaijr: <p>maybe somewhere here,</p>
<aside class="onebox githubblob">
  <header class="source">

      <a href="https://github.com/zdevito/ATen/blob/4aa3e1de29ed58457e530f84217e53db0998476c/aten/src/ATen/native/Normalization.cpp" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/zdevito/ATen/blob/4aa3e1de29ed58457e530f84217e53db0998476c/aten/src/ATen/native/Normalization.cpp" target="_blank" rel="noopener nofollow ugc">zdevito/ATen/blob/4aa3e1de29ed58457e530f84217e53db0998476c/aten/src/ATen/native/Normalization.cpp</a></h4>


    <pre><code class="lang-cpp">#include &lt;ATen/ATen.h&gt;
#include &lt;ATen/NativeFunctions.h&gt;
#include &lt;ATen/AccumulateType.h&gt;
#include &lt;ATen/CPUApplyUtils.h&gt;
#include &lt;ATen/Parallel.h&gt;
#include &lt;ATen/Config.h&gt;

#include &lt;ATen/detail/CUDAHooksInterface.h&gt;
#include &lt;ATen/native/TensorIterator.h&gt;
#include &lt;ATen/native/cpu/Loops.h&gt;

#include &lt;vector&gt;

static const int MIOPEN_DIM_MAX = 4;

namespace at { namespace native {

namespace {
  void check_dims_match_num_input_features(const char* arg_name, int64_t expected, int64_t actual){
    TORCH_CHECK(actual == expected,
</code></pre>


  This file has been truncated. <a href="https://github.com/zdevito/ATen/blob/4aa3e1de29ed58457e530f84217e53db0998476c/aten/src/ATen/native/Normalization.cpp" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>