<p dir="auto">I am working on adding batchnorm in the discriminator in WGAN-GP. However, I encountered a bug where <strong>gpu memory continues to increase</strong> when using <strong>batchnorm double backprop</strong>. This bug only occurs when using batchnorm. If i remove batchnorm from the model, the bug doesn't occur.</p>
<p dir="auto">Here's the code you can experiment with.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
import torch.nn as nn
from torch.autograd import Variable

# Model (partial discriminator)
D = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
    nn.BatchNorm2d(64),   # if you remove this, the bug does not occur.
    nn.LeakyReLU(0.2))

D.cuda()


for i in range(1000):
    # Input 
    x = Variable(torch.randn(10, 3, 128, 128).cuda(), requires_grad=True)
    out = D(x)
    
    grad = torch.autograd.grad(outputs=out,
                               inputs=x,
                               grad_outputs=torch.ones(out.size()).cuda(),
                               retain_graph=True,
                               create_graph=True,
                               only_inputs=True)[0]
    
    grad_norm = grad.pow(2).sum().sqrt()
    loss = torch.mean((grad_norm - 1)**2)

    # Reset grad and backprop
    D.zero_grad()
    loss.backward()
    
    if (i+1) % 10 == 0:
        print (i+1)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span> <span class="pl-k">import</span> <span class="pl-v">Variable</span>

<span class="pl-c"># Model (partial discriminator)</span>
<span class="pl-v">D</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">Conv2d</span>(<span class="pl-c1">3</span>, <span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span><span class="pl-c1">4</span>, <span class="pl-s1">stride</span><span class="pl-c1">=</span><span class="pl-c1">2</span>, <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-c1">1</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">BatchNorm2d</span>(<span class="pl-c1">64</span>),   <span class="pl-c"># if you remove this, the bug does not occur.</span>
    <span class="pl-s1">nn</span>.<span class="pl-v">LeakyReLU</span>(<span class="pl-c1">0.2</span>))

<span class="pl-v">D</span>.<span class="pl-en">cuda</span>()


<span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">1000</span>):
    <span class="pl-c"># Input </span>
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Variable</span>(<span class="pl-s1">torch</span>.<span class="pl-en">randn</span>(<span class="pl-c1">10</span>, <span class="pl-c1">3</span>, <span class="pl-c1">128</span>, <span class="pl-c1">128</span>).<span class="pl-en">cuda</span>(), <span class="pl-s1">requires_grad</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
    <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-v">D</span>(<span class="pl-s1">x</span>)
    
    <span class="pl-s1">grad</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">grad</span>(<span class="pl-s1">outputs</span><span class="pl-c1">=</span><span class="pl-s1">out</span>,
                               <span class="pl-s1">inputs</span><span class="pl-c1">=</span><span class="pl-s1">x</span>,
                               <span class="pl-s1">grad_outputs</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-s1">out</span>.<span class="pl-en">size</span>()).<span class="pl-en">cuda</span>(),
                               <span class="pl-s1">retain_graph</span><span class="pl-c1">=</span><span class="pl-c1">True</span>,
                               <span class="pl-s1">create_graph</span><span class="pl-c1">=</span><span class="pl-c1">True</span>,
                               <span class="pl-s1">only_inputs</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)[<span class="pl-c1">0</span>]
    
    <span class="pl-s1">grad_norm</span> <span class="pl-c1">=</span> <span class="pl-s1">grad</span>.<span class="pl-en">pow</span>(<span class="pl-c1">2</span>).<span class="pl-en">sum</span>().<span class="pl-en">sqrt</span>()
    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">mean</span>((<span class="pl-s1">grad_norm</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>)<span class="pl-c1">**</span><span class="pl-c1">2</span>)

    <span class="pl-c"># Reset grad and backprop</span>
    <span class="pl-v">D</span>.<span class="pl-en">zero_grad</span>()
    <span class="pl-s1">loss</span>.<span class="pl-en">backward</span>()
    
    <span class="pl-k">if</span> (<span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span>) <span class="pl-c1">%</span> <span class="pl-c1">10</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>:
        <span class="pl-en">print</span> (<span class="pl-s1">i</span><span class="pl-c1">+</span><span class="pl-c1">1</span>)</pre></div>
<p dir="auto">This is probably related to the previous <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="246937771" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/2264" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/2264/hovercard" href="https://github.com/pytorch/pytorch/issues/2264">#2264</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gchanan/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gchanan">@gchanan</a><br>
Running many iterations of double backward involving BatchNorm2d may cause out of memory error. There may be a memory leak somewhere? Code to reproduce the error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.autograd import Variable


class BatchNormTest(nn.Module):
    def __init__(self, c, num_classes=2):
        super(BatchNormTest, self).__init__()
        self.bn = nn.BatchNorm2d(c)

    def forward(self, x):
        out = x
        out = self.bn(out)
        out = F.relu(out)
        return out

c = 100
net = BatchNormTest(c)
use_cuda = True
inputs = Variable(torch.rand(100,c,100,100), requires_grad=True)
if use_cuda:
    net.cuda()
    inputs = inputs.cuda()

# on my server it fails at iteration 14
T = 100
for i in range(T):
    output = net(inputs)
    loss1 = torch.sum(output)
    grad_params = torch.autograd.grad(loss1, inputs, create_graph=True)

    grad = grad_params[0]
    loss = torch.sum(grad)

    loss.backward()
    print(i)"><pre class="notranslate"><code class="notranslate">import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.autograd import Variable


class BatchNormTest(nn.Module):
    def __init__(self, c, num_classes=2):
        super(BatchNormTest, self).__init__()
        self.bn = nn.BatchNorm2d(c)

    def forward(self, x):
        out = x
        out = self.bn(out)
        out = F.relu(out)
        return out

c = 100
net = BatchNormTest(c)
use_cuda = True
inputs = Variable(torch.rand(100,c,100,100), requires_grad=True)
if use_cuda:
    net.cuda()
    inputs = inputs.cuda()

# on my server it fails at iteration 14
T = 100
for i in range(T):
    output = net(inputs)
    loss1 = torch.sum(output)
    grad_params = torch.autograd.grad(loss1, inputs, create_graph=True)

    grad = grad_params[0]
    loss = torch.sum(grad)

    loss.backward()
    print(i)
</code></pre></div>
<p dir="auto">The error message:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="THCudaCheck FAIL file=/private/home/hongyizmit/pytorch/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory
Traceback (most recent call last):
  File &quot;models/testres.py&quot;, line 37, in &lt;module&gt;
    loss.backward()                                                                               
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py&quot;, line 156, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, retain_variables)
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/__init__.py&quot;, line 98, in backward
    variables, grad_variables, retain_graph)
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/nn/_functions/thnn/batchnorm_double_backwards.py&quot;, line 80, in batchnorm_double_backwards_fn
    gG = ggI * first_back_grad_input(gO, 1)
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/nn/_functions/thnn/batchnorm_double_backwards.py&quot;, line 73, in first_back_grad_input
    input_sub_mu.div(sigma2_eps) * sum_exclude_dim1(gO * input_sub_mu))
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py&quot;, line 829, in __mul__
    return self.mul(other)
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py&quot;, line 339, in mul
    return Mul.apply(self, other)
  File &quot;/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/_functions/basic_ops.py&quot;, line 48, in forward
    return a.mul(b)
RuntimeError: cuda runtime error (2) : out of memory at /private/home/hongyizmit/pytorch/torch/lib/THC/generic/THCStorage.cu:66"><pre class="notranslate"><code class="notranslate">THCudaCheck FAIL file=/private/home/hongyizmit/pytorch/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory
Traceback (most recent call last):
  File "models/testres.py", line 37, in &lt;module&gt;
    loss.backward()                                                                               
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py", line 156, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, retain_variables)
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/__init__.py", line 98, in backward
    variables, grad_variables, retain_graph)
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/nn/_functions/thnn/batchnorm_double_backwards.py", line 80, in batchnorm_double_backwards_fn
    gG = ggI * first_back_grad_input(gO, 1)
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/nn/_functions/thnn/batchnorm_double_backwards.py", line 73, in first_back_grad_input
    input_sub_mu.div(sigma2_eps) * sum_exclude_dim1(gO * input_sub_mu))
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py", line 829, in __mul__
    return self.mul(other)
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/variable.py", line 339, in mul
    return Mul.apply(self, other)
  File "/private/home/hongyizmit/.conda/envs/torchmaster/lib/python2.7/site-packages/torch/autograd/_functions/basic_ops.py", line 48, in forward
    return a.mul(b)
RuntimeError: cuda runtime error (2) : out of memory at /private/home/hongyizmit/pytorch/torch/lib/THC/generic/THCStorage.cu:66
</code></pre></div>