Tensors seem to be held by execution frames

NightMachinery: <p>My PyTorch code has a GPU memory leak.</p>
<p>To debug this issue, I used this function:</p>
<pre><code class="lang-auto">import gc
import torch
from pynight.common_torch import torch_memory_tensor


def find_tensors_on_gpu():
    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj):
                if obj.is_cuda:
                    obj_size = torch_memory_tensor(obj, s=2) #: MB
                    if obj_size &gt;= 10:
                        print(f'Tensor ID: {id(obj)} Type: {type(obj)} Size (MB): {obj_size}')
                        for ref in gc.get_referrers(obj):
                            try:
                                if isinstance(ref, dict):
                                    for k, v in ref.items():
                                        if v is obj:
                                            print(f'Variable Name: {k}')
                                else:
                                    print(f"ref: {ref}")
                            except Exception as e:
                                pass
        except Exception as e:
            pass

gc.collect()
find_tensors_on_gpu()
</code></pre>
<p>The output of this function is:</p>
<pre><code class="lang-auto">Tensor ID: 140152393953856 Type: &lt;class 'torch.Tensor'&gt; Size (MB): 23.0859375
Variable Name: features_out
ref: &lt;frame at 0x7f77c8dae640, file '/home/vit/code/pytorch-image-models/timm/models/decomposition.py', line 1873, code forward&gt;
Tensor ID: 140152393947616 Type: &lt;class 'torch.Tensor'&gt; Size (MB): 1142.75390625
Variable Name: attributions_v
ref: &lt;frame at 0x7f77c8dae640, file '/home/vit/code/pytorch-image-models/timm/models/decomposition.py', line 1873, code forward&gt;
Variable Name: attributions_v
Tensor ID: 140159443545328 Type: &lt;class 'torch.Tensor'&gt; Size (MB): 1142.75390625
Variable Name: residual_attributions_v
ref: &lt;frame at 0x55c5761c0270, file '/home/vit/code/pytorch-image-models/timm/models/vision_transformer.py', line 251, code forward&gt;
</code></pre>
<p>There are two problems:</p>
<ol>
<li>My GPU memory is at 7299MiB, while these tensors only sum up to about 3GB.</li>
<li>These tensors seem to be held by an execution frame of Python?! How do I free them?</li>
</ol>
<p>Related:</p>
<ul>
<li><a href="https://github.com/ipython/ipython/pull/11572" class="inline-onebox" rel="noopener nofollow ugc">fix a memory leak on exception (caused by the stored traceback) by stas00 · Pull Request #11572 · ipython/ipython · GitHub</a></li>
</ul>

ptrblck: <ol>
<li>PyTorch uses a caching allocator as described <a href="https://pytorch.org/docs/stable/notes/cuda.html#memory-management">here</a>.</li>
<li>The tensors seem to be referenced in the <code>forward</code> method so might be needed for the gradient computation in a <code>backward</code> call.</li>
</ol>

NightMachinery: <p>Is there a way to free up all the computational graph of a model’s parameters? Sth like <code>model.zero_grad()</code> but that frees up the stored computational graph.</p>
<p>PS: The Jupyter issue I linked was indeed my primary problem, and the workarounds there worked. In short, Jupyter stores references to local-scope tensors when an exception happens, and this seriously messes up the gc.</p>

ptrblck: <p>You could either call <code>backward()</code> on the output/loss which would free the intermediates or <code>del</code> the output of the model to delete all references to the intermediates and allow PyTorch to reuse the memory.</p>