Memory Leak because of dropout layer?

ash_gamma: <p>Training a MLP. [70 6000 6000 4] network clearly overfits on data. Tried adding dropout layers to prevent overfitting. Ran into <em>out of memory error</em>.</p>
<p>Line of thought as follows:</p>
<ol>
<li>
<p>Too many parameters, hence not fitting in 12Gb Tesla K80 GPU</p>
</li>
<li>
<p>Tried multi-GPU training by replacing <code>if cuda: model.cuda()</code> with <code>model=nn.DataParallel(model).cuda()</code>. Two 12Gb Tesla K80 GPUs run out of memory.</p>
</li>
<li>
<p>Tried training very small model (below) on 1 GPU. Still out ran out of memory. This network is smaller than even a MNIST MLP.</p>
</li>
</ol>
<p><strong>Specs:</strong></p>
<ol>
<li>Windows Server</li>
<li>Cuda 8</li>
<li>Python 3.6.1 (conda install)</li>
<li>PyTorch v0.4.0</li>
</ol>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.bn1 = nn.BatchNorm1d(70)
        self.fc1 = nn.Linear(70, 100)
        self.d1 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(100, 100)
        self.d2 = nn.Dropout(0.5)
        self.fc3 = nn.Linear(100, 4)

    def forward(self, x):
        x = x.view(-1, 70)
        x = self.bn1(x)
        x = F.relu(self.fc1(x))
        x = self.d1(x)
        x = F.relu(self.fc2(x))
        x = self.d2(x)
        return F.log_softmax(self.fc3(x))


model = Net()
if cuda:
    model.cuda()

optimizer = optim.Adam(model.parameters(), lr=0.00001, betas=(0.9, 0.999), weight_decay=0.0)
</code></pre>
<p>Checked discussions of few other reported cases of memory leaks.</p>
<ol>
<li>Added <code>torch.backends.cudnn.enabled = False</code> after all <code>import statements</code>. Didn’t help</li>
<li>
<code>volatile =True</code> during validation</li>
</ol>
<p>Any fixes?</p>

ptrblck: <p>Is your model with the Dropout working or are you getting the OOM error now all the time?<br>
Maybe a process still lives and uses all your GPU memory? Could you check it? I don’t know, if <code>nvidia-smi</code> works on a windows machine.</p>
<p>The <code>volatile</code> flag is deprecated. Since you are using pytorch <code>0.4.0</code>, you should use <code>with torch.no_grad()</code> instead.</p>
<p>Did you observe the memory? Is it growing once you add Dropout to your model?</p>

ash_gamma: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="18230">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Is your model with the Dropout working or are you getting the OOM error now all the time?</p>
</blockquote>
</aside>
<p>OOM when I add dropout layers. Without dropout, I could fit [70 6000 6000 4] model on 1 GPU. (Any model with dropout layer is running into OOM)</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="18230">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Maybe a process still lives and uses all your GPU memory? Could you check it? I don’t know, if <code>nvidia-smi</code> works on a windows machine.</p>
</blockquote>
</aside>
<p>I use <code>os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   os.environ["CUDA_VISIBLE_DEVICES"]="0,1"</code> to run it on specific GPUs. And no, these GPUs use 0 memory when experiments are not run. Checked via <code>nvidia-smi</code></p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="18230">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Did you observe the memory? Is it growing once you add Dropout to your model?</p>
</blockquote>
</aside>
<p>Yes, memory keeps growing when I try training models with dropout layer. I’ve trained bigger models (more parameters) on 1 12Gb GPU. But the model as small as one mentioned above, runs into OOM.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="18230">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>The <code>volatile</code> flag is deprecated. Since you are using pytorch <code>0.4.0</code> , you should use <code>with torch.no_grad()</code> instead.</p>
</blockquote>
</aside>
<p>Noted</p>

peterjc123: <p>I see that you are missing the if condition protection for the main process. Try again with that and see if that helps.</p>

ash_gamma: <aside class="quote no-group" data-username="peterjc123" data-post="4" data-topic="18230" data-full="true">
<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/peterjc123/48/2790_2.png" class="avatar"> peterjc123:</div>
<blockquote>
<p>I see that you are missing the if condition protection for the main process. Try again with that and see if that helps.</p>
</blockquote>
</aside>
<p>You mean like this: <code>if cuda: model=nn.DataParallel(model).cuda()</code> ?</p>
<p>Tried. Didn’t help</p>

peterjc123: <p>No, I mean you might protect your code with an idiom on Windows, like</p>
<pre><code class="lang-python">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.bn1 = nn.BatchNorm1d(70)
        self.fc1 = nn.Linear(70, 100)
        self.d1 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(100, 100)
        self.d2 = nn.Dropout(0.5)
        self.fc3 = nn.Linear(100, 4)

    def forward(self, x):
        x = x.view(-1, 70)
        x = self.bn1(x)
        x = F.relu(self.fc1(x))
        x = self.d1(x)
        x = F.relu(self.fc2(x))
        x = self.d2(x)
        return F.log_softmax(self.fc3(x))

if __name__ == '__main__':
    model = Net()
    if cuda:
        model.cuda()

    optimizer = optim.Adam(model.parameters(), lr=0.00001, betas=(0.9, 0.999), weight_decay=0.0)
    # and more code that is on the outer part without a protection
</code></pre>

ash_gamma: <p>Memory usage increases at every iteration, for every epoch. (checked via <code>nvidia-smi</code>)</p>
<p>This is my <em>main</em> function now. I’m running it from jupyter notebook.</p>
<p>Tried with and without having data loaders in <em>main</em></p>
<pre><code class="lang-auto">if __name__ == '__main__':
    
    train_data = torch.utils.data.TensorDataset(features_t,labels_t)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size, shuffle=True)

    val_data = torch.utils.data.TensorDataset(features_v,labels_v)
    validation_loader = torch.utils.data.DataLoader(val_data, batch_size, shuffle=False)
    
    model = Net()
    if cuda:
        model = nn.DataParallel(model).cuda()

    optimizer = optim.Adam(model.parameters(), lr=0.00001,betas=(0.9, 0.999), weight_decay=0.0)
    
    %%time
    epochs = 6000

    losst, acct = [], []
    lossv, accv = [], []
    best_acc = 0.0
    for epoch in range(1, epochs + 1):

        train(epoch, losst, acct)
        a = validate(lossv, accv)
        print("Mean Class Acc" + str(a))

        is_best = a &gt; best_acc
        best_acc = max(a, best_acc)
        save_checkpoint({
            'epoch': epoch + 1,
            'net':model,
            'state_dict': model.state_dict(),
            'best_prec1': best_acc,
            'optimizer' : optimizer.state_dict(),
        }, is_best)
</code></pre>

peterjc123: <p><a class="mention" href="/u/ash_gamma">@ash_gamma</a> Could you please also post the functions of <code>train</code> and <code>validate</code> ? In 0.4.0, you should remember to use <code>with torch.no_grad():</code> during interference.</p>

ash_gamma: <pre><code class="lang-auto">def train(epoch, loss_vector_train, accuracy_vector_train ,log_interval=1000):
    model.train()
    train_loss, correct = 0, 0
    for batch_idx, (data, target) in enumerate(train_loader):
        if cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        
        loss = F.nll_loss(output, target, weight=weights.cuda())
        train_loss += loss
        pred = output.data.max(1)[1]
        
        correct += pred.eq(target.data).cpu().sum()
        loss.backward()
        optimizer.step()
        
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            epoch, batch_idx * len(data), len(train_loader.dataset),
            100. * batch_idx / len(train_loader), loss.data[0]))
            
        
    train_loss /= len(train_loader)
    loss_vector_train.append(train_loss)
    accuracy = 100. * correct / len(train_loader.dataset)
    accuracy_vector_train.append(accuracy)
    print(str(train_loss) + " " + str(accuracy))

</code></pre>
<pre><code class="lang-auto">def validate(loss_vector, accuracy_vector):
    model.eval()
    val_loss, correct = 0, 0
    confusion_matrix = torchnet.meter.ConfusionMeter(4)
    with torch.no_grad():
        for data, target in validation_loader:
            if cuda:
                data, target = data.cuda(), target.cuda()
            data, target = Variable(data), Variable(target)
            output = model(data)
            confusion_matrix.add(output.data, target.data)
            val_loss += F.nll_loss(output, target, weight=weights.cuda()).data[0]
            pred = output.data.max(1)[1] 
            correct += pred.eq(target.data).cpu().sum()

    val_loss /= len(validation_loader)
    loss_vector.append(val_loss)

    accuracy = 100. * correct / len(validation_loader.dataset)
    accuracy_vector.append(accuracy)
    cm = confusion_matrix.conf
    d = np.diag(cm)
    s = np.sum(cm,1)
    mean_class_acc = np.mean(d/s)
    print('\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        val_loss, correct, len(validation_loader.dataset), accuracy))
    
    return mean_class_acc
</code></pre>

peterjc123: <p>The use of <code>train_loss += loss</code> may be the cause. Try using <code>train_loss += loss.item()</code> instead.<br>
Please refer to the latest <a href="https://github.com/pytorch/examples/blob/master/mnist/main.py" rel="nofollow noopener">example</a> of MNIST to change your code.</p>

ash_gamma: <p>I’ll try these changes.</p>
<p>But to let you know, I’m able to run training scripts without running into OOM when I do not have dropout layer. Shouldn’t <code>train_loss +=</code> be a problem even without dropout then?</p>

peterjc123: <p>I don’t think the problem is on the dropout side. As it is used so many examples in PyTorch repo, users should have reported them since it’s now almost one month after the last release. And before the release, we have ran several benchmarks on various networks, which contains AlexNet that has dropout in it.</p>

ash_gamma: <aside class="quote no-group" data-username="peterjc123" data-post="10" data-topic="18230">
<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/peterjc123/48/2790_2.png" class="avatar"> peterjc123:</div>
<blockquote>
<p>The use of <code>train_loss += loss</code> may be the cause. Try using <code>train_loss += loss.item()</code></p>
</blockquote>
</aside>
<p>This was the issue! Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Cerbal: <p>Interestingly, the same thing happened to me today.</p>
<p>The memory leak was also caused by a bad ‘+=’ between a float and a scalar tensor.<br>
But the problem only appeared when I added a dropout layer.</p>

OaDsis: <p>same issue for me, and can be fixed by <a href="https://discuss.pytorch.org/u/peterjc123">peterjc123</a>’s method,it is really wierd:)</p>