Repository: pytorch/pytorch

Issue:
JIT leaks memory in nets with dropout layer 
if I add a dropout layer in FC layers, a JIT compiled forward pass of a net starts to seriously leak memory in the backward pass. I made a simple test case below. When it runs, you will see a rapidly increasing GPU/CPU memory usage in nvidia-smi/top and it will crash with an out of memory exception soon. Without the dropout layer or the backward pass, it works fine. 

```python
import torch
from torch import jit
import torch.nn as nn
from torch.autograd import Variable

class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.net1 = nn.Linear(100, 200)
        self.net2 = nn.Linear(200, 1)
        self.sigmoid = nn.Sigmoid()
        self.ReLU = nn.ReLU(inplace=False)
        self.drop = nn.Dropout(0.5)
               
    def forward(self, V):
        return self.sigmoid(self.net2(self.drop(self.ReLU(self.net1(V))))).squeeze() 


use_cuda = True
net = TestNet()
criterion = nn.BCELoss()
if use_cuda:
    net.cuda()
    criterion.cuda()
    V = Variable(torch.randn(100, 100)).cuda()
    label = Variable(torch.randn(100)).cuda()
else:
    V = Variable(torch.randn(100, 100))
    label = Variable(torch.randn(100))

net.train()
fwd = jit.compile(net.forward)
for i in range(0,1000000):
    r = fwd(V)
    err = criterion(r, label)
    err.backward()   

```




Discussion:
Thanks for the report! Note that the JIT is still unreleased, but of course feel free to play with it. The proper way to use it with nn Modules is to decorate the whole class, but it's still worthwhile to investigate this leak
I use somewhat complex multi-head network with lots of conditionals in the forward function. The conditionals do not change during the training of the corresponding parts of the network, so I could unroll all of the conditionals into multiple classes, but that would create a lot of classes and clutter in the code, making it less readable and manageable.  So, I am hoping that I could JIT compile the forward function on demand, possibly multiple times during different phases of training when the conditionals are known to be static? So, I guess decorating the whole class will not work in my case.

Every class instance has its own trace cache, so it should still work for you
If I decorate the whole class, it seems to work without a noticeable speedup. So, I don't know if it uses a compiled structure or not. Can it be related to a fact that I also call the forward method from the class __init__ method to dynamically create the network structure?

I like that in pytorch I can build the network structure dynamically, so from the "__init__" I call the "forward" while setting a boolean flag "create" to create the network structure based on a structure of the first sample of the data. In all of the subsequent "forward" calls, the network structure does not change nor depend on conditionals. Can the JIT compiler ignore this first forward invocation from the __init__ method to avoid creating the precompiled graph? 





Hmm no I don't think we will be supporting something like this already. An implementation detail is that calling `.forward` will probably go around the JIT, and only `module(...)` will trigger it, so you might try. But as I said, it can break any moment.

Also, I wouldn't expect a noticeable speedup on this kind of model. It's quite small and GEMM-heavy, so there's not a lot of things we can optimize yet
I was talking about a real model that has hundreds of python lines in the forward call. It executes considerably faster (50-70% faster) if compiled with the jit.compile method, but unfortunately leaks the memory when the dropout is included. With the decorations I see no speedup. The simple code above is just as a greatly simplified test case for the memory leak.

It would be nice if jit.compile would work, as the decorations seem to be too restrictive when building more complex dynamic networks.


That's great to hear. We'll look into the dropout issue. You can try using `model.jit_debug_info()` and see the number of JIT cache hits/misses 
Could you provide us with a dummy script that showcases your use case? I'm not sure if I fully understand what you're doing, so it would be helpful to see the code
I may also be experiencing a related leak, though I haven't really triaged yet; I get an OOM after a few dozen training iterations when I annotate the first of these modules (inside a larger network; the second module shows the direct context where it's used):
```
@torch.jit.compile(nderivs=1)
class LayerNorm(nn.Module):

    def __init__(self, d_model, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

class ResidualBlock(nn.Module):

    def __init__(self, layer, d_model, drop_ratio, pos=0):
        super(ResidualBlock, self).__init__()
        self.layer = layer
        self.dropout = nn.Dropout(drop_ratio)
        self.layernorm = LayerNorm(d_model)
        self.pos = pos

    def forward(self, *x):
        return self.layernorm(x[self.pos] + self.dropout(self.layer(*x)))
```
I also have to explicitly specify the `super` call Python 2-style; if I use `super()` I get:
```
Traceback (most recent call last):
  File "/opt/conda/envs/pytorch-py3.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 171, in __init__
    old_init(self, *args, **kwargs)
  File "/home/james/transformer/model.py", line 191, in __init__
    super().__init__()
TypeError: super(type, obj): obj must be an instance or subtype of type

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "run.py", line 464, in <module>
    model = args.model(SRC, TRG, args)
  File "/home/james/transformer/model.py", line 858, in __init__
    self.encoder = Encoder(src, args)
  File "/home/james/transformer/model.py", line 534, in __init__
    [EncoderLayer(args) for i in range(args.n_layers)])
  File "/home/james/transformer/model.py", line 534, in <listcomp>
    [EncoderLayer(args) for i in range(args.n_layers)])
  File "/home/james/transformer/model.py", line 464, in __init__
    args.d_model, args.drop_ratio)
  File "/home/james/transformer/model.py", line 207, in __init__
    self.layernorm = LayerNorm(d_model)
  File "/opt/conda/envs/pytorch-py3.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 177, in __init__
    "\n\nOriginal error: {}".format(str(e))), e)
  File "<string>", line 3, in raise_from
TypeError: torch.jit.compile must be used as a class decorator; using it on an already defined class is not valid.

Original error: super(type, obj): obj must be an instance or subtype of type
```
I'll look into that today

Resolving patch:
diff --git a/torch/nn/_functions/dropout.py b/torch/nn/_functions/dropout.py
--- a/torch/nn/_functions/dropout.py
+++ b/torch/nn/_functions/dropout.py
@@ -25,20 +25,22 @@ def forward(cls, ctx, input, p=0.5, train=False, inplace=False):
         ctx.train = train
         ctx.inplace = inplace
 
+        if ctx.p == 0 or not ctx.train:
+            return input
+
         if ctx.inplace:
             ctx.mark_dirty(input)
             output = input
         else:
             output = input.clone()
 
-        if ctx.p > 0 and ctx.train:
-            ctx.noise = cls._make_noise(input)
-            if ctx.p == 1:
-                ctx.noise.fill_(0)
-            else:
-                ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)
-            ctx.noise = ctx.noise.expand_as(input)
-            output.mul_(ctx.noise)
+        ctx.noise = cls._make_noise(input)
+        if ctx.p == 1:
+            ctx.noise.fill_(0)
+        else:
+            ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)
+        ctx.noise = ctx.noise.expand_as(input)
+        output.mul_(ctx.noise)
 
         return output
 


Test patch: