GPU memory not fully released after training loop

username: abweiss
cooked: <p>My training + cross-validation loop usually looks something like:</p>
<pre><code class="lang-python">for epoch in range(max_epoch):
    print('Epoch:',epoch)

    # TRAINING
    mean_err = 0
    net.train()
    for train_idx, train_data in enumerate(train_dataloader):
        # wrap train_data in Variables; move to GPU
        # zero gradients
        # do forward/backward propagation stuff
        # update mean_err
    print('Mean Training Error:', mean_err)
    
    # CROSS-VALIDATION
    mean_err = 0
    net.eval()
    for cv_idx, cv_data in enumerate(cv_dataloader):
        # wrap cv_data in Variables; move to GPU
        # do forward propagation stuff
        # update mean_err
    print('Mean Cross-Validation Error:', mean_err)
</code></pre>
<p>What I often find is that there is a particular batch size in which the training loop runs just fine, but then the GPU immediately runs out of memory and crashes in the cross-validation loop.  If I step down the batch size, the system becomes stable, and will run indefinitely through many epochs.  What’s going on?  The inference-only cross-validation should be less memory intensive than the training, so obviously not all the memory used during the training loop is being freed.  How do I fix this so that I can run larger batches?</p>

username: Bobak_Farzin
cooked: <p>I see the same problem.  The train works fine, but as I try to run validation data the memory on the card increases till I get a CUDA out of memory error.  Very interested in a solution to this.</p>
<p>FWIW, my model is an RNN with LSTM units.  Not sure if that matters.</p>

username: spruceb
cooked: <p>My guess is in the training loop you assign various variables, likely something like <code>prediction = net(train_data)</code>. This <code>prediction</code> has pointers to the rest of the entire computation graph the data goes through (so probably a lot of memory). Once the training loop finishes, the last assignment to <code>prediction</code> is still kept, since in Python loops and <code>if</code>s don’t create new scopes. All those variables are in the function scope (or global scope). So there’s no way to know that that whole graph should now be reclaimed.</p>
<p>Either explicitly <code>del</code> the variables after the training loop, or put the training loop into a separate function.</p>
<p>By they way, also make sure you’re not doing something like <code>mean_err += err</code> where <code>err</code> is a <code>Variable</code>, because that will also create a large unnecessary graph. Use <code>err.data</code> instead.</p>

username: Bobak_Farzin
cooked: <p>I am already using the <code>err.data[0]</code>, but that is good to know.  <code>del</code> of variables does not seem to free up the CUDA memory at all. Is there something I need to do for CUDA vars after I call <code>del</code>?</p>
<p>For my use case, the difference between <code>train</code> and <code>validate</code> is the calls to:</p>
<pre><code>opt.zero_grad()
...
loss.backward() 
opt.step()
</code></pre>
<p>Could those be freeing up the memory in some way?  My train loop has far more batches and has no issues with GPU memory.</p>

username: spruceb
cooked: <p>I don’t think there’s anything special you need to do with CUDA variables, I haven’t in the past.<br>
Just to be sure, when you say it doesn’t free up memory do you mean you still crash, or are you using a profiler of some kind? So far as I know the memory will still be claimed by PyTorch for later use, so still used from any profiler’s perspective.</p>
<p>To check if it’s memory unfreed from training or actually something wrong with your validation loop, you can try commenting out the training loop entirely and use the larger batch size.</p>

username: Bobak_Farzin
cooked: <p>My <code>nvidia-smi</code> goes to the max memory on the card and the loop will crash with:</p>
<p><code>THCudaCheck FAIL file=/b/wheel/pytorch-src/torch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory</code></p>

username: spruceb
cooked: <p>Are you deleting the direct output of the <code>net</code> and any other variables calculated from that output, say your <code>loss</code>?</p>

username: Bobak_Farzin
cooked: <p>Per your ealier point, the graph is building on itself and continuing to increase in size. It was subtle detail for me as a newbie to PyTorch.</p>
<p>I am using an RNN with LSTM units.  I have to detach the hidden states at each pass otherwise I build up the size of the network.</p>
<p>From an example, I had already gotten:<br>
<code>def repackage_hidden(h): """Wraps hidden states in new Variables, to detach them from their history.""" if type(h) == Variable: return Variable(h.data) else: return tuple(repackage_hidden(v) for v in h)</code><br>
but I was not calling that on my validation loop.  Adding that in, I can proceed and validation works with no trouble.</p>

username: JiamingSun
cooked: <p>After debugging for hours I found this line of code solved my problem:</p>
<pre><code class="lang-python">with torch.no_grad():
    valid()  # function of your validation loop
</code></pre>
<p>Strongly suggest this should be put into official tutorials and examples.</p>

username: Mata_Fu
cooked: <p><a class="mention" href="/u/jiamingsun">@JiamingSun</a><br>
Could you write an example? I have met the same problem and I can’t understand your solution.<br>
Thanks a lot!</p>

username: JiamingSun
cooked: <pre><code class="lang-auto">for epoch in range(num_epochs):
    train()  # function contains the training loop.
    with torch.no_grad():
        valid()  # function contains the validation loop.
</code></pre>
<p>I think I’ve made it more clear this time.</p>

username: xin_Wang1
cooked: <p>I also have same problem in my code. Luckily, I found the method to solve it.<br>
you can set torch.backends.cudnn.benchmark=False</p>