Uncollected object references from torch.load() cause memory growth

TShimko126: <p>Hi all,</p>
<p>I’m currently working on training a model on a rather large dataset ( ~1 billion examples in total). To store and load this dataset, I’m employing the Webdataset library (<code>https://github.com/webdataset/webdataset</code>). I am storing preprocessed examples as Pytorch tensors in a tar file, per the Webdataset spec. I am running into a memory issue after a long period of training where the machine performing the reading/loading operation runs out of memory. After a bit of debugging, it seems to be caused by some unexpected (at least to me) behavior with the Pytorch serialization code.</p>
<p>Since each tensor is stored as a separate file, then tarred together, each must be loaded individually, meaning that ~1 billion torch.load() operations take place per epoch. I traced memory usage using the <code>tracemalloc</code> utility and discovered that some part of this serialization process is creating objects that are not garbage collected even after the tensor read from disk goes out of scope (or in the example below, is manually deleted).</p>
<p>I have included a minimal example below to show these objects. The example creates a dummy tensor, then reads it into memory repeatedly. The memory usage is queried in 10,000 step intervals. Given that that the tensor goes out of scope immediately after it is read with torch.load(), I would not expect to see any objects from <code>torch/serialization.py</code> to still be in scope. However, I instead see ~10,000 new objects created each time the memory is queried (every 10,000 steps, see example output below). These objects also seem to never be released and cause constant linear memory growth.</p>
<p>Is there something I am missing here? Is there a way to remove references to these objects so they can be collected? Even though they are only 28B each, on average, they still seem to be causing memory use to grow linearly until the machine is out of memory (usually after many millions of examples are read from disk on a high-mem training machine).</p>
<p>Thank you in advance for any and all help you can offer!</p>
<p>Best,<br>
Tyler</p>
<h2>
<a name="system-info-1" class="anchor" href="#system-info-1"></a>System info</h2>
<ul>
<li>Ubuntu 18.04.1</li>
<li>Python 3.8</li>
<li>Pytorch 1.9.0</li>
</ul>
<h2>
<a name="minimal-reproducible-example-2" class="anchor" href="#minimal-reproducible-example-2"></a>Minimal reproducible example:</h2>
<pre><code class="lang-auto">import torch
import tracemalloc

if __name__ == '__main__':
    # Create a dummy Tensor serialized to disk
    rand = torch.rand((20, 10))
    torch.save(rand, 'test.pth')

    # Start tracemalloc
    tracemalloc.start(30)
    old_snapshot = tracemalloc.take_snapshot()

    for i in range(100000):
        # Load the Tensor,
        test_tensor = torch.load('test.pth')

        # Immediately delete reference to test_tensor
        del test_tensor

        if i % 10000 == 0 and i != 0:
            # Take snapshot
            snapshot = tracemalloc.take_snapshot()

            # Print changes in memory consumption
            print(f'################# STEP {i} #################')
            for stat in snapshot.compare_to(old_snapshot, 'lineno')[:2]:
                print(str(stat))
            print('############################################')

            # Save snapshot
            old_snapshot = snapshot

</code></pre>
<h2>
<a name="example-output-3" class="anchor" href="#example-output-3"></a>Example output:</h2>
<pre><code class="lang-auto">################# STEP 10000 #################
.../python3.8/site-packages/torch/serialization.py:845: size=274 KiB (+274 KiB), count=10005 (+10005), average=28 B
.../python3.8/site-packages/torch/serialization.py:242: size=274 KiB (+274 KiB), count=10003 (+10003), average=28 B
################# STEP 20000 #################
.../python3.8/site-packages/torch/serialization.py:845: size=547 KiB (+273 KiB), count=20005 (+10000), average=28 B
.../python3.8/site-packages/torch/serialization.py:242: size=547 KiB (+273 KiB), count=20003 (+10000), average=28 B
################# STEP 30000 #################
.../python3.8/site-packages/torch/serialization.py:242: size=820 KiB (+273 KiB), count=30003 (+10000), average=28 B
.../python3.8/site-packages/torch/serialization.py:845: size=821 KiB (+273 KiB), count=30004 (+9999), average=28 B
</code></pre>
<p>The two lines being referenced from <code>torch/serialization.py</code> above are:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L845">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L845" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L845" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L845</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="835" style="counter-reset: li-counter 834 ;">
          <li>
          </li>
<li>def _load(zip_file, map_location, pickle_module, pickle_file='data.pkl', **pickle_load_args):</li>
          <li>    restore_location = _get_restore_location(map_location)</li>
          <li>
          </li>
<li>    loaded_storages = {}</li>
          <li>
          </li>
<li>    def load_tensor(data_type, size, key, location):</li>
          <li>        name = f'data/{key}'</li>
          <li>        dtype = data_type(0).dtype</li>
          <li>
          </li>
<li class="selected">        storage = zip_file.get_storage_from_record(name, size, dtype).storage()</li>
          <li>        loaded_storages[key] = restore_location(storage, location)</li>
          <li>
          </li>
<li>    def persistent_load(saved_id):</li>
          <li>        assert isinstance(saved_id, tuple)</li>
          <li>        typename = _maybe_decode_ascii(saved_id[0])</li>
          <li>        data = saved_id[1:]</li>
          <li>
          </li>
<li>        assert typename == 'storage', \</li>
          <li>            f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'"</li>
          <li>        data_type, key, location, size = data</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>and</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L242">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L242" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L242" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L242</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="232" style="counter-reset: li-counter 231 ;">
          <li>        if 'w' in mode:</li>
          <li>            return _open_buffer_writer(name_or_buffer)</li>
          <li>        elif 'r' in mode:</li>
          <li>            return _open_buffer_reader(name_or_buffer)</li>
          <li>        else:</li>
          <li>            raise RuntimeError(f"Expected 'r' or 'w' in mode but got {mode}")</li>
          <li>
          </li>
<li>
          </li>
<li>class _open_zipfile_reader(_opener):</li>
          <li>    def __init__(self, name_or_buffer) -&gt; None:</li>
          <li class="selected">        super(_open_zipfile_reader, self).__init__(torch._C.PyTorchFileReader(name_or_buffer))</li>
          <li>
          </li>
<li>
          </li>
<li>class _open_zipfile_writer_file(_opener):</li>
          <li>    def __init__(self, name) -&gt; None:</li>
          <li>        super(_open_zipfile_writer_file, self).__init__(torch._C.PyTorchFileWriter(str(name)))</li>
          <li>
          </li>
<li>    def __exit__(self, *args) -&gt; None:</li>
          <li>        self.file_like.write_end_of_file()</li>
          <li>
          </li>
<li>
      </li>
</ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


VitalyFedyunin: <p><code>del x</code> doesn’t guarantee garbage collection in Python, you need to try with <code>gc.collect()</code></p>

TShimko126: <p>Hi <a class="mention" href="/u/vitalyfedyunin">@VitalyFedyunin</a>, thanks for the suggestion! Unfortunately, the issue persists even after adding an explicit call to <code>gc.collect()</code> after each read. I’ve included an updated script and output below.</p>
<pre><code class="lang-python">import torch
import tracemalloc
import gc

if __name__ == '__main__':
    # Create a dummy Tensor serialized to disk
    rand = torch.rand((20, 10))
    torch.save(rand, 'test.pth')

    # Start tracemalloc
    tracemalloc.start(30)
    old_snapshot = tracemalloc.take_snapshot()

    for i in range(30001):
        # Load the Tensor
        test_tensor = torch.load('test.pth')

        # Immediately delete reference to test_tensor
        del test_tensor
        gc.collect()

        if i % 10000 == 0 and i != 0:
            # Take snapshot
            snapshot = tracemalloc.take_snapshot()

            # Print changes in memory consumption
            print(f'################# STEP {i} #################')
            for stat in snapshot.compare_to(old_snapshot, 'lineno')[:2]:
                print(str(stat))
            print('############################################')

            # Save snapshot
            old_snapshot = snapshot
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">################# STEP 10000 #################
.../python3.8/site-packages/torch/serialization.py:845: size=274 KiB (+274 KiB), count=10003 (+10003), average=28 B
.../python3.8/site-packages/torch/serialization.py:242: size=274 KiB (+274 KiB), count=10003 (+10003), average=28 B
############################################
################# STEP 20000 #################
.../python3.8/site-packages/torch/serialization.py:845: size=547 KiB (+273 KiB), count=20003 (+10000), average=28 B
.../python3.8/site-packages/torch/serialization.py:242: size=547 KiB (+273 KiB), count=20003 (+10000), average=28 B
############################################
################# STEP 30000 #################
.../python3.8/site-packages/torch/serialization.py:845: size=820 KiB (+273 KiB), count=30003 (+10000), average=28 B
.../python3.8/site-packages/torch/serialization.py:242: size=820 KiB (+273 KiB), count=30003 (+10000), average=28 B
############################################
</code></pre>
<p>I’ve also opened an issue on the Pytorch Github here:</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/64043">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/64043" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/64043" target="_blank" rel="noopener nofollow ugc">Uncleared memory use after torch.load()</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-08-26" data-time="18:58:01" data-timezone="UTC">06:58PM - 26 Aug 21 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/TShimko126" target="_blank" rel="noopener nofollow ugc">
          <img alt="TShimko126" src="https://avatars.githubusercontent.com/u/3413296?v=4" class="onebox-avatar-inline" width="20" height="20">
          TShimko126
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: serialization
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🐛 Bug

When loading a tensor from disk using `torch.load()`, the full memor<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">y used by the loading process does not get collected even when the tensor goes out of scope. The amount of uncleared memory grows with the number of `torch.load()` operations.

## To Reproduce

Steps to reproduce the behavior:

Memory usage is measured using [guppy3](https://pypi.org/project/guppy3/).

1. Run the following script to get baseline memory usage after creating and deleting a random tensor:

```python
from torch import rand
from guppy import hpy

# Get the heap before Tensor creation
hp = hpy()
before = hp.heap()

# Create a dummy Tensor
rand = rand((20, 10))
del rand

# Get the after heap and diff
after = hp.heap()
leftover = after - before
print(leftover)
```

Output:
```
Partition of a set of 1 object. Total size = 408 bytes.
 Index  Count   %     Size   % Cumulative  % Kind (class / dict of class)
     0      1 100      408 100       408 100 types.FrameType
```

2. Compare to the following output which creates a tensor, saves it to disk, then reads it with `torch.load()` and deletes it:

```python
from torch import rand, save, load
from guppy import hpy

# Create a dummy Tensor and save to disk
rand = rand((20, 10))
save(rand, 'test.pth')
del rand

# Get the heap before torch.load()
hp = hpy()
before = hp.heap()

# Read and del
test_tensor = load('test.pth')
del test_tensor

# Get the after heap and diff
after = hp.heap()
leftover = after - before
print(leftover)
print('\nGet referrer to leftover dict:')
print(leftover[1].byid.referrers.byvia)
```

Output:
```
Partition of a set of 2 objects. Total size = 640 bytes.
 Index  Count   %     Size   % Cumulative  % Kind (class / dict of class)
     0      1  50      408  64       408  64 types.FrameType
     1      1  50      232  36       640 100 dict (no owner)

Get referrer to leftover dict:
Partition of a set of 1 object. Total size = 416 bytes.
 Index  Count   %     Size   % Cumulative  % Referred Via:
     0      1 100      416 100       416 100 "['Unpickler']", "['_unpickler']", '.__objclass__',
                                             '.__self__', '[0]'
```

Given that the uncleared `dict` appears to be referred to by the `Unpickler` object, my thought is that this may be related to the wrapper defined [here](https://github.com/pytorch/pytorch/blob/master/torch/serialization.py#L869), but I may be completely wrong.

3. Finally, run this script to show that the memory usage allocated to the `dict` grows with subsequent `torch.load()` operations:

```python
from torch import rand, save, load
from guppy import hpy

# Create a dummy Tensor and save to disk
rand = rand((20, 10))
save(rand, 'test.pth')
del rand

# Get the heap before torch.load()
hp = hpy()
before = hp.heap()

# Read and del
for i in range(1000):
    test_tensor = load('test.pth')
    del test_tensor
del i

# Get the after heap and diff
after = hp.heap()
leftover = after - before
print(leftover)
print('\nGet referrer to leftover dict:')
print(leftover[0].byid.referrers.byvia)
```

Output:
```
Partition of a set of 2 objects. Total size = 9720 bytes.
 Index  Count   %     Size   % Cumulative  % Kind (class / dict of class)
     0      1  50     9312  96      9312  96 dict (no owner)
     1      1  50      408   4      9720 100 types.FrameType
```

Note the size increase of the `dict` from 408 -&gt; 9312 between the examples in step 2 and step 3.

## Expected behavior

All memory occupied during/after the `torch.load()` operation should be cleared after the tensor is deleted/goes out of scope.

## Environment

```
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 10.15.7 (x86_64)
GCC version: Could not collect
Clang version: 12.0.0 (clang-1200.0.32.27)
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.5 (default, Sep  4 2020, 02:22:02)  [Clang 10.0.0 ] (64-bit runtime)
Python platform: macOS-10.15.7-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.4
[pip3] pytorch-lightning==1.2.10
[pip3] torch==1.9.0
[pip3] torch-cluster==1.5.9
[pip3] torch-geometric==1.7.1
[pip3] torch-scatter==2.0.7
[pip3] torch-sparse==0.6.10
[pip3] torch-spline-conv==1.2.1
[pip3] torchmetrics==0.2.0
[conda] blas                      1.0                         mkl
[conda] mkl                       2019.5                intel_281    intel
[conda] mkl-service               2.3.0            py38h9ed2024_0
[conda] mkl_fft                   1.2.0            py38hc64f4ea_0
[conda] mkl_random                1.1.1            py38h959d312_0
[conda] numpy                     1.19.4                   pypi_0    pypi
[conda] numpy-base                1.19.2           py38hcfb5961_0
[conda] pytorch-lightning         1.2.10                   pypi_0    pypi
[conda] torch                     1.9.0                    pypi_0    pypi
[conda] torch-cluster             1.5.9                    pypi_0    pypi
[conda] torch-geometric           1.7.1                    pypi_0    pypi
[conda] torch-scatter             2.0.7                    pypi_0    pypi
[conda] torch-sparse              0.6.10                   pypi_0    pypi
[conda] torch-spline-conv         1.2.1                    pypi_0    pypi
[conda] torchmetrics              0.2.0                    pypi_0    pypi
```

## Additional context

I previously brought this issue up on the [Pytorch forum](https://discuss.pytorch.org/t/uncollected-object-references-from-torch-load-cause-memory-growth/130188) and will include the relevant text below:

I’m currently working on training a model on a rather large dataset ( ~1 billion examples in total, with multiple individual tensors that need to be loaded per each example). To store and load this dataset, I’m employing the [Webdataset library](https://github.com/webdataset/webdataset). I am storing preprocessed examples as Pytorch tensors in a tar file, per the Webdataset spec. I am running into a memory issue after a long period of training where the machine performing the reading/loading operation runs out of memory. After a bit of debugging, it seems to be caused by some unexpected (at least to me) behavior with the Pytorch serialization code.

Since each tensor is stored as a separate file, then tarred together, each must be loaded individually, meaning that ~1 billion torch.load() operations take place per epoch. I traced memory usage using the tracemalloc utility and discovered that some part of this serialization process is creating objects that are not garbage collected even after the tensor read from disk goes out of scope (or in the example below, is manually deleted).

I have included a minimal example below to show these objects. The example creates a dummy tensor, then reads it into memory repeatedly. The memory usage is queried in 10,000 step intervals. Given that that the tensor goes out of scope immediately after it is read with torch.load(), I would not expect to see any objects from torch/serialization.py to still be in scope. However, I instead see ~10,000 new objects created each time the memory is queried (every 10,000 steps, see example output below). These objects also seem to never be released and cause constant linear memory growth. Even though they are only 28B each, on average, they still seem to be causing memory use to grow linearly until the machine is out of memory (usually after many millions of examples are read from disk on a high-mem training machine).

Thank you in advance for any and all help you can offer! I sincerely apologize if this has been addressed elsewhere and I missed it while searching. Right now this is blocking bug for one of my ongoing projects.

```python
import torch
import tracemalloc

if __name__ == '__main__':
    # Create a dummy Tensor serialized to disk
    rand = torch.rand((20, 10))
    torch.save(rand, 'test.pth')

    # Start tracemalloc
    tracemalloc.start(30)
    old_snapshot = tracemalloc.take_snapshot()

    for i in range(30001):
        # Load the Tensor
        test_tensor = torch.load('test.pth')

        # Immediately delete reference to test_tensor
        del test_tensor

        if i % 10000 == 0 and i != 0:
            # Take snapshot
            snapshot = tracemalloc.take_snapshot()

            # Print changes in memory consumption
            print(f'################# STEP {i} #################')
            for stat in snapshot.compare_to(old_snapshot, 'lineno')[:2]:
                print(str(stat))
            print('############################################')

            # Save snapshot
            old_snapshot = snapshot

```

Output:
```
################# STEP 10000 #################
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:845: size=274 KiB (+274 KiB), count=10004 (+10004), average=28 B
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:242: size=274 KiB (+274 KiB), count=10003 (+10003), average=28 B
############################################
################# STEP 20000 #################
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:845: size=547 KiB (+273 KiB), count=20005 (+10001), average=28 B
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:242: size=547 KiB (+273 KiB), count=20003 (+10000), average=28 B
############################################
################# STEP 30000 #################
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:242: size=820 KiB (+273 KiB), count=30003 (+10000), average=28 B
/usr/local/Caskroom/miniconda/base/envs/graphene/lib/python3.8/site-packages/torch/serialization.py:845: size=821 KiB (+273 KiB), count=30004 (+9999), average=28 B
############################################
```

The relevant referenced lines are:
- [torch/serialization.py:845](https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L845)
- [torch/serialization.py:242](https://github.com/pytorch/pytorch/blob/v1.9.0/torch/serialization.py#L242)



cc @mruberry</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>