A clear explanation of what num_workers=0 means for a DataLoader

Machupicchu: <p>Hello,</p>
<p>the pytorch documentation  it says that setting num_workers=0 for a DataLoader causes it to be handled by the “main process” from the pytorch doc:</p>
<p>" <code>0</code> means that the data will be loaded in the main process."</p>
<p>maybe i’m wrong but usually i find that the pytorch doc gives often (but not always of course) many useless or obvious info but does not mention the only useful points that i m looking for…and its quite frustrating some times.</p>
<p>… by reading a bit more in other places that seems to mean no multiprocessing i.e. single-process. ok. But does it mean that the whole dataset (all the minibatches) are loaded into the main memory (RAM) ?</p>
<p>can you confirm that this is the case (this is actually what i want because i have a dataset of around 1 million files) which is extremely slow for even a 16-worker dataloader to load and i can afford having the whole dataset (around 32 GB) loaded into RAM so the minibatches would be fetched quite fast after that.</p>
<p>From what i read elsewhere it seems to be the case 0 meaning loaded into RAM but it would be great if the pytorch experts of this forum can confirm</p>
<p>P.S.<br>
the pin_memory=True and setting e.g. a prefetch_factor=32 or something does not help, so ideally and since i can afford it : load the whole dataset in RAM but “transparently” i.e. if i can just set num_workers=0 that would be the dream.</p>

ptrblck: <aside class="quote no-group" data-username="Machupicchu" data-post="1" data-topic="177614">
<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/m/b19c9b/48.png" class="avatar"> Machupicchu:</div>
<blockquote>
<p>But does it mean that the whole dataset (all the minibatches) are loaded into the main memory (RAM) ?</p>
</blockquote>
</aside>
<p>The <code>DataLoader</code> is not deciding how and when the actual samples are loaded, which is defined in the <code>Dataset</code>. Depending how the <code>Dataset.__init__</code> and <code>.__getitem__</code> are implemented the entire internal data could be preloaded or each sample could be lazily loaded.</p>
<p>The <code>DataLoader</code> will call into <code>Dataset.__getitem__</code> to load each sample and will create the batch using the <code>collate_fn</code>. Each worker will create a copy of the <code>Dataset</code> (assuming you are using <code>num_workers&gt;0</code>).</p>
<aside class="quote no-group" data-username="Machupicchu" data-post="1" data-topic="177614">
<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/m/b19c9b/48.png" class="avatar"> Machupicchu:</div>
<blockquote>
<p>and i can afford having the whole dataset (around 32 GB) loaded into RAM so the minibatches would be fetched quite fast after that.</p>
</blockquote>
</aside>
<p>This is defined by the <code>Dataset</code> as mentioned above.</p>
<aside class="quote no-group" data-username="Machupicchu" data-post="1" data-topic="177614">
<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/m/b19c9b/48.png" class="avatar"> Machupicchu:</div>
<blockquote>
<p>From what i read elsewhere it seems to be the case 0 meaning loaded into RAM but it would be great if the pytorch experts of this forum can confirm</p>
</blockquote>
</aside>
<p>That’s not the case.</p>
<aside class="quote no-group" data-username="Machupicchu" data-post="1" data-topic="177614">
<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/m/b19c9b/48.png" class="avatar"> Machupicchu:</div>
<blockquote>
<p>maybe i’m wrong but usually i find that the pytorch doc gives often (but not always of course) many useless or obvious info but does not mention the only useful points that i m looking for…and its quite frustrating some times.</p>
</blockquote>
</aside>
<p>Improvements to the docs are always welcome, so feel free to submit fixes or missing information.</p>

Machupicchu: <p>oh thanks for the very fast answer. I did implement the <strong>init</strong> and <strong>getitem</strong> and i m actually listing all the files from a folder and then the getitem returns a minibatch that i form from the contents of the file…<br>
i was under the impresseion (but misled by a certain GPT!) that setting num_workers=0 would “magically” i.e. internally cause the dataloader to pre_load the whole dataset in RAM</p>
<p>chatGPT pretends that:</p>
<pre><code class="lang-auto">from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

train_data = MNIST('data/', train=True, download=True, transform=ToTensor())
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers=0)

# All minibatches are loaded into RAM before training starts
for batch in train_loader:
    # Do something with the batch
    pass
</code></pre>
<p>so GPT is saying nonsense (again) ? XD</p>
<p>(i could load it all manually it’s true, but each file is 1MB so id need to open it , do some preproc extract the data etc… i.e. do what is done in the getitem sort of… and save it in a variable i guess… but that would require some more code to write and i hoped but maybe just in my (chatGPT’s) dreams that num_workers=0 would do it)</p>

ptrblck: <p>The comment is true, since the <code>MNIST</code> dataset will preload all samples as seen <a href="https://github.com/pytorch/vision/blob/b78d98bb152ffb9c0c0f5365f59f475c70b1784e/torchvision/datasets/mnist.py#L104">here</a>. The <code>__getitem__</code> will then index this preloaded data <a href="https://github.com/pytorch/vision/blob/b78d98bb152ffb9c0c0f5365f59f475c70b1784e/torchvision/datasets/mnist.py#L138">here</a> and process it afterwards.<br>
You can access the preloaded data also via:</p>
<pre><code class="lang-python">train_data.data.shape
# torch.Size([60000, 28, 28])
</code></pre>
<p>However, this behavior is not defined by the <code>DataLoader</code> or its number of workers, but by the <code>Dataset</code> (<code>MNIST</code> in this case).</p>

Machupicchu: <p>oh i see of course with the dummy MNIST example… i have my own train_data that i wrote of course</p>
<p>the funny thing is that for e.g. 350k files the perf is only slow for the first epoch and then the data seems to be in RAM and everything is much faster the next epoch… (this is the case even when not specifying pin_memory or prefetch_factor) however when going to aroung 700k it remains be slow even after the 0,1,2 epochs … do you have an idea why? is it like there is a “threshold” around 700k that is just too much to handle maybe? I was under the impression that after the first epoch if all the minibatches can stay in RAM, then they should remain there, which seems to be the case with 350k but not 700k</p>

ptrblck: <aside class="quote no-group" data-username="Machupicchu" data-post="5" data-topic="177614">
<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/m/b19c9b/48.png" class="avatar"> Machupicchu:</div>
<blockquote>
<p>I was under the impression that after the first epoch if all the minibatches can stay in RAM, then they should remain there, which seems to be the case with 350k but not 700k</p>
</blockquote>
</aside>
<p>I also don’t know where this understanding comes from but the <code>DataLoader</code> is also not caching the data in any way. It would still be your responsibility to write a caching mechanism into your custom <code>Dataset</code>.</p>

Machupicchu: <p>well I said that because I noticed it on my computer: first epoch → very slow, 2nd-last epochs → much faster… so clearly some type of caching was happening … but what triggered it ?</p>

Machupicchu: <p>I mean of course I am sure you are right about the DataLoader so the “sort of caching mechanism” that i did notice must come from somewhere else… do you think it could be a “smart feature” of the OS (Ubuntu in my case) ?!</p>

ptrblck: <p>Yes, your OS could also cache the data locally but I would also guess the cache might still be valid after the initial run (i.e. even if you rerun the same Python script).<br>
To clear the cache you could run <code>/sbin/sysctl vm.drop_caches=3</code> between each Python script execution or epoch.</p>