Overall memory usage using multi-process data loading with Subset

octave: <p>Hello,</p>
<p>From the warning note in the multi-process data loading <a href="https://pytorch.org/docs/stable/data.html#multi-process-data-loading" rel="noopener nofollow ugc">multi-process data loading section</a>, the overall memory usage when using multiple workers is <code>number of workers * size of parent process</code>.</p>
<p>When using torch.utils.data.Subset, does <code>size of parent process</code> correspond to the original dataset or to a subset of the original dataset?</p>
<p>In other words, when splitting a Dataset in <code>number of subsets</code> Subset and loading them using DataLoader with <code>number workers</code> each, is the overall memory usage <code>number of workers * size of parent process</code>.or <code>number of subsets * number of workers * size of parent process</code>?</p>

ptrblck: <p>Each worker will create a copy of the <code>Dataset</code>. The size of the <code>Dataset</code> is defined by the <code>__init__</code> method and if you are lazily loading each sample (i.e. if you are only storing the paths in the <code>__init__</code>), the additional memory usage would be tiny. The <code>Subset</code> will only wrap the <code>Dataset</code> and pass the provided indices to it. It won’t change the actual <code>Dataset.__init__</code> method and the lazy/eager loading logic will still be applied.</p>

octave: <p>Thank you for your answer.</p>
<p>So the overall memory would indeed scale linearly with the <code>number of subsets</code> but with lazy loading the size of the original <code>Dataset</code> can be shrunk at the cost of the time to read from disk at each <code>Dataset.__getitem__</code> call?</p>

ptrblck: <aside class="quote no-group" data-username="octave" data-post="3" data-topic="186585">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/f05b48/48.png" class="avatar"> octave:</div>
<blockquote>
<p>So the overall memory would indeed scale linearly with the <code>number of subsets</code></p>
</blockquote>
</aside>
<p>No, since a <code>Subset</code> is only wrapping a <code>Dataset</code> and passes indices to it. The memory increase would come from storing the additional indices.</p>
<aside class="quote no-group" data-username="octave" data-post="3" data-topic="186585">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/f05b48/48.png" class="avatar"> octave:</div>
<blockquote>
<p>but with lazy loading the size of the original <code>Dataset</code> can be shrunk at the cost of the time to read from disk at each <code>Dataset.__getitem__</code> call?</p>
</blockquote>
</aside>
<p>Yes, lazily loading the data will only use memory to store the paths etc. while initializing the dataset.</p>

octave: <p>Do you have a script, or ideas on how to write one, to monitor the memory of all the data loading processes for educational purpose?</p>
<p>I tried adapting <a href="https://ppwwyyxx.com/blog/2022/Demystify-RAM-Usage-in-Multiprocess-DataLoader/" rel="noopener nofollow ugc">this blog article</a> and <a href="https://github.com/ppwwyyxx/RAM-multiprocess-dataloader" rel="noopener nofollow ugc">code</a>:</p>
<pre><code class="lang-auto">import torch
from torch.utils.data import Dataset, Subset, DataLoader
from common import MemoryMonitor

class ToyDataset(Dataset):
  def __init__(self, shape):
    self.data = torch.zeros(shape)

  def __len__(self):
    return self.data.shape[0]

  def __getitem__(self, idx):
    return self.data[idx]

shape = ... # tuple, (num samples, height, width)
dataset = ToyDataset(shape)

dataloader = DataLoader(dataset, batch_size=32, num_workers=4, persistent_workers=True)
it = iter(dataloader)
monitor = MemoryMonitor()
[monitor.add_pid(w.pid) for w in it._workers]

print(f"Single Dataset, {dataloader.num_workers} workers\n", monitor.table())
</code></pre>
<p>From the blog above:</p>
<blockquote>
<p>By definition, we should <strong>use total PSS to count the total RAM usage</strong> of N processes.</p>
</blockquote>
<p>For the same number of worker the memory used by each worker scales with the size of the data as expected.</p>
<p>I’m not sure how to interpret the outputs when keeping the size of the dataset constant and varying the number of workers e.g. from <code>num_workers=2</code> and <code>num_workers=4</code>:</p>
<pre><code class="lang-auto"># For a dataset of shape (512, 1024, 1024) and batch size 32

Single Dataset, 2 workers
   time    PID  rss    pss     uss     shared    shared_file
------  -----  -----  ------  ------  --------  -------------
 21097  98405  2.4G   930.6M  157.1M  2.3G      43.4M
 21097  98598  2.4G   889.7M  119.4M  2.3G      31.0M
 21097  98600  2.4G   906.9M  136.3M  2.3G      31.0M
</code></pre>
<pre><code class="lang-auto"># For a dataset of shape (512, 1024, 1024) and batch size 32

Single Dataset, 4 workers
   time     PID  rss    pss     uss     shared    shared_file
------  ------  -----  ------  ------  --------  -------------
 21996  102287  2.4G   620.7M  158.5M  2.2G      43.4M
 21996  102488  2.4G   583.5M  126.2M  2.2G      31.3M
 21996  102490  2.4G   598.1M  140.6M  2.2G      31.5M
 21996  102492  2.3G   547.9M  90.3M   2.2G      31.6M
 21996  102494  2.3G   566.1M  108.3M  2.2G      31.6M
</code></pre>
<p><strong>The memory usage per worker seems to decrease with increasing number of workers.</strong> Do you know if this is expected or if it’s a issue with my example?</p>