Free all GPU memory used in between runs

JSLJ23: <p>Hi pytorch community,</p>
<p>I was hoping to get some help on ways to completely free GPU memory after a single iteration of model training.</p>
<p>This process is part of a Bayesian optimisation loop involving a molecular docking program that runs on the GPU as well so I cannot terminate the code halfway to “free” the memory.<br>
The cycle looks something like this:</p>
<ol>
<li>Run docking</li>
<li>Train model to emulate docking</li>
<li>Run inference and choose the best data points</li>
<li>Repeat (10 times or so)</li>
</ol>
<p>In between each step of docking and model training, pytorch seems to hold on to a block of memory as depicted in <code>nvtop</code> and <code>nvidia-smi</code> and despite me deleting the <code>model</code>, and <code>optimizer</code> by calling <code>del</code> on them, as well as running <code>gc.collect()</code> &amp; <code>torch.cuda.empty_cache()</code>, this memory is still being taken up and my docking program runs into OOM errors. The model training and inference are nicely wrapped into a function and don’t exceed that function’s scope as the predictions are converted to numpy arrays and appended to a pandas DataFrame. I’ve tried my best to isolate the ML stuff from the docking stuff as much as possible.</p>
<p>I was hoping to get some help on this and if there are robust ways to free GPU memory fully, I would definitely be keen to include them.</p>
<p>Thank you!</p>

ptrblck: <p>Deleting all objects and references pointing to objects allocating GPU memory is the right approach and will free the memory. Calling <code>empty_cache()</code> will also clear the cache and free the memory (besides the memory used for the CUDA context).<br>
Here is a small example:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn


def memory_stats():
    print(torch.cuda.memory_allocated()/1024**2)
    print(torch.cuda.memory_cached()/1024**2)


def allocate():
    x = torch.randn(1024*1024, device='cuda')
    memory_stats()
    
    
memory_stats()
# 0.0
# 0.0

allocate()
# 4.0 # allocated inside the function
# 20.0 # used cache

memory_stats()
# 0.0 # local tensor is free
# 20.0 # cache is still alive

torch.cuda.empty_cache()
memory_stats()
# 0.0
# 0.0 # cache is free again

x = torch.randn(1024, 1024, device='cuda')
memory_stats()
# 4.0
# 20.0

# store referece
y = x

del x # this does not free the memory of x since y still points to it
memory_stats() 
# 4.0  
# 20.0

del y # this allows PyTorch to free the memory and reuse it in the cache
memory_stats()
# 0.0
# 20.0

torch.cuda.empty_cache()
memory_stats() 
# 0.0
# 0.0
</code></pre>

JSLJ23: <p>Thanks for the quick reply.</p>
<p>Is there a systematic way to perform these delete operations for the actual <code>model</code>, <code>optimizer</code> and <code>data</code> moved onto the GPU from enumerating the <code>data_loader</code> during after training &amp; inference?</p>
<p>For a simple single <code>Tensor</code> example I tired, it works fine and makes sense, but I can’t seem to figure out exactly where to place them into the training loop to free the memory.</p>

ptrblck: <p>You can <code>del</code> any object you’ve created, but would need to make sure no references are stored to it as my example code shows.<br>
You can replace the actual <code>x</code> tensor in my code with a real model etc. and would see the same effect.<br>
Users often claim to see “memory leaks” by e.g. storing references to computation graphs and other tensors via logging information e.g. for the loss:</p>
<pre><code class="lang-python">output = model(input)
loss = criterion(output, target)
losses += loss # !!! This will store the entire computation graph !!!
</code></pre>
<p>so you might want to check if you are running into similar issues.</p>

jvwilliams23: <p>Hi, I was wondering, is there an attribute or function to help find potential unfreed references to a tensor? Similar to how <code>grad_fn</code> can help understand the upstream graph for graph issues.</p>
<p>I see many people say e.g. you need to make sure no references to the variable are stored, but it is unclear how to do that in big complex codebases. Any ideas? Thanks!</p>