Does DataLoader preload future batches? If yes, how to turn it off?

StrausMG: <p>Hi! I am working on a simple classification problem. However, in my setup, I would like to create batches smarter than just by uniform sampling. Namely, I am trying to mine hard batches as following:</p>
<ol>
<li>sample a big batch uniformly (e.g. 1024 samples)</li>
<li>apply my model to the big batch and calculate losses</li>
<li>sample a normal batch (e.g. 128 samples) out of the big batch using multinomial distribution parameterized by the losses from step 2</li>
</ol>
<p>This procedure depends on the model and the model changes after every batch. Consequently, I do not want to preload future batches. In other words, I want the method  <code>__next__</code> of <code>dataloader.batch_sampler</code> to be called only when the same method of <code>dataloader</code> itself is explicitly called. For example, in the following loop I would like the next batch to be created only after <code>do</code> returns</p>
<pre><code class="lang-python">for batch in dataloader:
    do(batch)
</code></pre>
<p>I know that I can achieve that by setting <code>num_workers=0</code>. But what is the behavior if I set <code>num_workers=some_positive_integer</code> ? Does <code>Dataloader</code> preload future batches? If yes, is it possible to avoid it anyhow?</p>
<p>P.S. My final goal is to create my custom BatchSampler and the question, if answered, will help to understand if it is possible to yield batches instead of precalculating them for the whole epoch.</p>

ptrblck: <p>If you set <code>num_workers &gt; 0</code>, each worker will load a batch in the background and inplace modifications of your <code>Dataset</code> won’t be noticed until the next epoch.<br>
As you said, <code>num_workers=0</code> might work and I can’t see any reason to use multiple workers, as your workflow seems to be sequential, i.e. you can just sample the next batch after <code>do</code> finishes.<br>
So even if you would use multiple workers, all would have to wait until your loop finishes.</p>

StrausMG: <p>Thank you for the response! I was not sure about batch-level parallelization, I thought that <code>num_workers &gt; 0</code> may boost even construction of one batch. Now I understand that it works in a different way.</p>

Logan_Nielsen: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> could I get a little more info on how this works? I’m assuming <code>__getitem__</code> eventually gets called on the dataset so as long as what gets returned there is the desired result then preprocessing is occurring.</p>
<p>One other question, the preparation step in my application is extensive and unavoidable. I’m thinking about using multiple threads to help the word go faster. Do I need to worry about how these processes affect what Pytorch is doing in preparing the data? (e.g. Pytorch thinks a thread isn’t being used when really I’ve just kicked up a process on it.)</p>

ptrblck: <aside class="quote no-group" data-username="Logan_Nielsen" data-post="4" data-topic="38760">
<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/logan_nielsen/48/43945_2.png" class="avatar"> Logan_Nielsen:</div>
<blockquote>
<p>I’m assuming <code>__getitem__</code> eventually gets called on the dataset so as long as what gets returned there is the desired result then preprocessing is occurring.</p>
</blockquote>
</aside>
<p>Usually you would add the proprocessing in the <code>__getitem__</code> and would return the already preprocessed batch of samples. Note that this thread had a specific requirement:</p>
<blockquote>
<p>Consequently, I do not want to preload future batches. In other words, I want the method <code>__next__</code> of <code>dataloader.batch_sampler</code> to be called only when the same method of <code>dataloader</code> itself is explicitly called.</p>
</blockquote>
<p>So I’m unsure if you are working with a similar requirement or if it’s a general question.</p>
<aside class="quote no-group" data-username="Logan_Nielsen" data-post="4" data-topic="38760">
<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/logan_nielsen/48/43945_2.png" class="avatar"> Logan_Nielsen:</div>
<blockquote>
<p>I’m thinking about using multiple threads to help the word go faster. Do I need to worry about how these processes affect what Pytorch is doing in preparing the data? (e.g. Pytorch thinks a thread isn’t being used when really I’ve just kicked up a process on it.)</p>
</blockquote>
</aside>
<p>You should be able to use multi-threading inside a process, i.e. in each worker of your <code>DataLoader</code>. Could you describe your concern a bit more, please? I.e. are you seeing any specific issues in your current use case?</p>