DataLoader relationship between num_workers, prefetch_factor and type of DataSet

username: pumplerod
cooked: <p>I’m slowly wrapping my head around proper use of a custom DataSet as used by the DataLoader.</p>
<p>I have a <code>torch.utils.data.Dataset</code> which utilizes a:</p>
<pre><code class="lang-auto">def __getitem__(self, idx):
    ...
    return { "image": X_ims, "label": Y}
</code></pre>
<p>when I supply the DataLoader with a batch_size I get the expected number of samples back.  But the documentation is a bit unclear with regard to what the <code>workers</code> are doing.  The documentation mentions that they are getting samples from the DataSet, and I am “assuming” that the <code>prefetch_factor</code> is the number of samples they are retrieving.  Do they do this until another entire batch is loaded ( ideally while my model is training)?  or is each worker trying to collect a full batch of samples?</p>
<p>The question arises because I find I am often stuck waiting on DataLoader’s<br>
<code>_next_data() &gt; _get_data() &gt; _try_get_data() &gt; _get() &gt; _wait()</code> method.  Which I take to mean that it is waiting for the workers to get back with all the samples.</p>
<p>If I have a large batch size ( 1000 ), assuming I have <code>num_workers=2</code> should I set the <code>prefetch_factor=500</code>?  In the hope each worker is able to get enough samples for the next batch by the time I’m ready for it?</p>

username: ptrblck
cooked: <p>Each worker will load a complete batch in the background (if <code>num_workers&gt;0</code>) while the training is executed. There is a feature request to allow multiple workers to work on a single batch, but I don’t think it’s implemented yet.</p>
<p>The <code>prefetch_factor</code> defines the number of batches, which are preloaded, if I’m not mistaken, so 500 would be quite large (it could be alright, if you have enough memory).</p>

username: agrpz
cooked: <p>I see that the recent version rather seems to load prefetch_factor*num_workers “samples”.</p>
<p>prefetch_factor (int, optional, keyword-only arg): Number of samples loaded<br>
in advance by each worker. <code>2</code> means there will be a total of<br>
2 * num_workers samples prefetched across all workers. (default: <code>2</code>)</p>
<p>However, using different prefetch_factor values did not absolutely change the used GPU memory <a href="https://github.com/aabrol/SMLvsDL/blob/f3362cf6bfab6e479983ebb09bfd87554a8e62b4/reprex/utils.py#L343" rel="noopener nofollow ugc">for my pipeline</a>. But not sure if it is due to the customized dataloader or another issue with this newer pytorch functionality (hoping to spend more time on this soon, but would appreciate any feedback if someone happens to stop by to look at this).</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="agrpz" data-post="3" data-topic="117735">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agrpz/48/37313_2.png" class="avatar"> agrpz:</div>
<blockquote>
<p>However, using different prefetch_factor values did not absolutely change the used GPU memory</p>
</blockquote>
</aside>
<p>I wouldn’t expect it to change the GPU memory usage, since the data is preloaded on the host and each batch is pushed to the device in the training loop in the common use case. This also seems to be your workflow, as I cannot see any GPU usage in the <code> MRIDataset</code>.</p>

username: agrpz
cooked: <p>Thanks much! I was assuming this functionality enables data to be prefetched to GPU (via CPU), but looks like from your explanation it is limited to CPU memory (and that is rather what I must monitor for noting performance gains). Although this functionality would be helpful, I guess it would still result in a squarish waveform for GPU utilization (idle times). I have actually tried calling cuda within <strong>getitem</strong> in the past with little success (I remember the code threw errors but might try again).</p>