Simultaneously preprocess a batch on CPU and run forward/backward on GPU

sammlapp: <p>Hi, this is a basic understanding question: how can I make a DataLoader start preparing the next batch while a model runs .forward() and .backward() on the previous batch?</p>
<p>When we write a loop to train or run inference with a PyTorch model, we loop over a DataLoader to get the next batch of samples, then run model.forward(x) within the loop (and if training also calculate loss and run model.backward()).</p>
<p>Simple inference loop:</p>
<pre><code class="lang-python">for batch in dataloader: #loads next batch, which happens on CPU and can take a while.
    batch_tensors = batch["X"].to(device_name)
    logits = model.forward(batch_tensors)  #runs forward pass on GPU
    # the next batch doesn't start loading until this finishes (I assume)
</code></pre>
<p>In a scenario where the DataLoader uses several CPU workers to parallelize data, and the model.forward() and model.backward() are on the GPU, it seems like the workload will always bet trading off between the CPU (while DataLoader gets the next batch) and GPU (while model.forward runs, for instance). In my use case, pre-processing with the DataLoader can be heavy such that the time to prepare a batch is equal to the time to run the forward pass. So, it seems like if the DataLoader would immediately start preparing the next batch instead of waiting for the loop to complete, inference would be twice as fast.</p>
<p>Am I missing something here? Perhaps the DataLoader does actually start working on the next batch immediately rather than waiting for the top of the loop to be reached again.</p>
<p>If not, is there a way to achieve this simultaneous use of CPU preprocessing and GPU forward/backward passes?</p>
<p>Thank you for your help, I understand this may be an obvious or simple question for more experienced users.</p>

nivek: <p>You should be able to do that with the argument <code>prefetch_factor</code> for <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader" rel="noopener nofollow ugc"><code>DataLoader</code> (documentation)</a>:</p>
<ul>
<li>
<strong>prefetch_factor</strong> (<a href="https://docs.python.org/3/library/functions.html#int" rel="noopener nofollow ugc"><em>int</em></a><em>,</em> <em>optional</em> <em>,</em> <em>keyword-only arg</em>) – Number of batches loaded in advance by each worker. <code>2</code> means there will be a total of 2 * num_workers batches prefetched across all workers. (default: <code>2</code>)</li>
</ul>
<p>Let us know if you are seeing unexpected behaviors.</p>

ptrblck: <aside class="quote no-group" data-username="sammlapp" data-post="1" data-topic="161987">
<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/sammlapp/48/36795_2.png" class="avatar"> sammlapp:</div>
<blockquote>
<p>Am I missing something here? Perhaps the DataLoader does actually start working on the next batch immediately rather than waiting for the top of the loop to be reached again.</p>
</blockquote>
</aside>
<p>Yes, if multiple workers are used they will prepare the next batch and won’t wait until the training loop finished. Once the queue is full, the workers will wait as <a class="mention" href="/u/nivek">@nivek</a> described and you can change the behavior via the <code>prefetch_factor</code>.<br>
Here is a small example:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.arange(100).view(-1, 1)
        
    def __getitem__(self, idx):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is not None:
            print("worker {} loading index {}".format(worker_info.id, idx))
        x = self.data[idx]
        return x
    
    def __len__(self):
        return len(self.data)

dataset = MyDataset()
loader = DataLoader(dataset, num_workers=2)

for data in loader:
    print("executing training")
    time.sleep(5)
    print("training done")
</code></pre>
<p>Output:</p>
<pre data-code-wrap="python"><code class="lang-python">worker 0 loading index 0
worker 1 loading index 1
worker 0 loading index 2
worker 1 loading index 3
worker 0 loading index 4
executing training
worker 1 loading index 5
training done
executing training
worker 0 loading index 6
training done
executing training
worker 1 loading index 7
training done
executing training
worker 0 loading index 8
training done
executing training
worker 1 loading index 9
training done
executing training
...
</code></pre>
<p>The data loading is fast in this case as it’s simple indexing and you can see that the workers directly fill the queue before the “training” (which is a <code>sleep</code> call here) can even start. Once the data is consumed, the next worker will start loading and processing the next batch.</p>