Queue in DataLoader does not behave as expected when using num_workers?

lkc1: <p>Do I understand the following correctly?</p>
<p>When <code>num_workers &gt;=1</code>, the main process pre-loads <code>prefetch_factor * num_workers</code> batches. When the training loop consumes one batch, the corresponding worker loads the next batch in its queue.</p>
<p>If this is the case, let’s go through an example.</p>
<p><em>NOTE: I have chosen the numeric values for illustration purposes and</em><br>
<em>have ignored various overheads in this example. The accompanying code example uses these numbers.</em></p>
<p>Say I have <code>num_workers=4</code>,  <code>prefetch_factor=4</code>,  <code>batch_size=128</code>. Further assume, it takes <code>0.003125 s</code> to fetch an item from a source database and the train step takes <code>0.05 s</code>.</p>
<p>Now, each batch would take <code>0.003125 * 128 = 0.4 s</code> to load.</p>
<p>With a <code>prefetch_factor=4</code> and <code>num_workers=4</code>, first, <code>4*4=16</code> batches will be loaded.</p>
<p>Once the <code>16</code> batches are loaded, the first train step consumes 1 batch and takes <code>0.05 s</code>. Say <code>worker[0]</code> provided this batch and will start the process to generate a new batch to replenish the queue. Recall fetching a new batch takes <code>0.4 s</code>.</p>
<p>Similarly, the second step consumes one more batch and the corresponding worker (<code>worker[1]</code> in this example) starts the data fetching process.</p>
<p>The first 8 train steps would take <code>0.05*8=0.4s</code>. By this time, <code>8</code> batches have been<br>
consumed and <code>worker[0]</code> has produced 1 batch. In the next step, 1 batch is consumed and <code>worker[1]</code> produces a new batch. <code>worker[1]</code> had started the data fetching process in the second train step which would now be completed.</p>
<p>Following this we can see, each subsequent train step will consume 1 batch and one of the workers will produce 1 batch, keeping the dataloader queue to have always <code>8</code> batches. This means that the train step is never waiting for the data loading process as there are always 8 batches in the buffer.</p>
<p>I would expect this behavior regardless of the data size of the batch given <code>num_workers</code>, <code>prefetch_factor</code> are large enough. However, in the following code example that is not case.</p>
<p>In the code below, I define a custom iterable that returns a numpy array. As the size of the numpy array increases, increasing <code>num_worker</code> or ‘prefetch_factor’ does not improve the time taken for running through a batch.</p>
<p>I’m guessing this is because each worker serializes the batch to send to the main process where it is de-serialized. As the data size increase, this process would take more time. However, I would think if the queue size is large enough (<code>num_workers</code>, <code>prefetch_factor</code>), at some point, there should be a break even point where each training step consumption of a batch would be accompanied by replenishment via one of the workers as I illustrated in the above example.</p>
<p>In the code below, when <code>MyIterable</code> returns a small object (np array of size (10, 150)), increasing <code>num_workers</code> helps as expected. But when the returned object is larger (np array of size (1000, 150)), <code>num_workers</code> or <code>prefetch_factor</code> does not do much.</p>
<pre><code class="lang-auto"># small np object
avg time per batch for num workers=0: 0.47068126868714444
avg time per batch for num workers=2: 0.20982365206225495
avg time per batch for num workers=4: 0.10560789656221914
avg time per batch for num workers=6: 0.07202646931250456
avg time per batch for num workers=8: 0.05311137337469063
</code></pre>
<pre><code class="lang-auto"># large np object
avg time per batch for num workers=0: 0.6090951558124971
avg time per batch for num workers=2: 0.4594530961876444
avg time per batch for num workers=4: 0.45023533212543043
avg time per batch for num workers=6: 0.3830978863124983
avg time per batch for num workers=8: 0.3811495694375253
</code></pre>
<p>Am I missing something here? Why doesn’t the data loader queue have enough buffer such that data loading is not the bottleneck?</p>
<p>Even if the serialization and de-serialization process would take longer for the latter case, I’d expect to have a large enough buffer where the consumption and replenishment rate of the batches are almost equal. Otherwise, what is the point of having <code>prefetch_factor</code>.</p>
<p>If the code is behaving as expected, are there any other ways to pre-load the next <code>n</code> batches in a buffer such that it is large enough and never depleted?</p>
<p>Thanks</p>
<pre><code class="lang-auto">import time
import torch
import numpy as np
from time import sleep
from torch.utils.data import DataLoader, IterableDataset


def collate_fn(records):
    # some custom collation function
    return records


class MyIterable(object):
    def __init__(self, n):
        self.n = n
        self.i = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.i &lt; self.n:
            sleep(0.003125)  # simulates data fetch time
            # return np.random.random((10, 150))  # small data item
            return np.random.random((1000, 150))  # large data item
        else:
            raise StopIteration


class MyIterableDataset(IterableDataset):
    def __init__(self, n):
        super(MyIterableDataset).__init__()
        self.n = n

    def __iter__(self):
        return MyIterable(self.n)


def get_performance_metrics(num_workers):
    ds = MyIterableDataset(n=10000)

    if num_workers == 0:
        dl = torch.utils.data.DataLoader(ds, num_workers=0, batch_size=128, collate_fn=collate_fn)
    else:
        dl = torch.utils.data.DataLoader(ds, num_workers=num_workers, prefetch_factor=4, persistent_workers=True,
                                         batch_size=128, collate_fn=collate_fn,
                                         multiprocessing_context='spawn')
    warmup = 5
    times = []
    t0 = time.perf_counter()
    for i, batch in enumerate(dl):
        sleep(0.05)  # simulates train step
        e = time.perf_counter()
        if i &gt;= warmup:
            times.append(e - t0)
        t0 = time.perf_counter()
        if i &gt;= 20:
            break
    print(f'avg time per batch for num workers={num_workers}: {sum(times) / len(times)}')


if __name__ == '__main__':
    num_worker_options = [0, 2, 4, 6, 8]
    for n in num_worker_options:
        get_performance_metrics(n)
</code></pre>

Arij-Aladel: <p>I have tested it also I think the answer is here<br>
where as written in the documentation <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p><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>)</p>
<p>but there is try in the emplementation this indicates that it will iterate until 2 * num_workers so this is themaximum number if the core is busy then the total number will be less than that<br>
according to the function _try_put_index  in the <a href="https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html#DataLoader" rel="noopener nofollow ugc">implementation</a></p>

Vinayaka_Hegde: <p><a class="mention" href="/u/arij-aladel">@Arij-Aladel</a> I dont follow, could you please elaborate? <code>_try_put_index</code> doesnot have any logic for iterating until <code>2*no_workers</code> as you mentioned.</p>

Vinayaka_Hegde: <p><a class="mention" href="/u/lkc1">@lkc1</a> where you able to figure this out ?</p>

lkc1: <p>So the problem isn’t Pytorch, it is with how multiprocessing works in python which Pytorch uses. In python multiprocessing, all the data objects in the parent process are serialized to be passed to the child process. If the size of the parent objects is not large, the overhead is not a problem and you see the advantage of using multiple processes in total time. However, if the data objects in the parent process are large (for my case it was ~100-200 MB), the serialization process adds overhead and thereby neutralizing any benefits one might get with multiprocessing.</p>
<p>So, I ended up with a somewhat complicated solution of creating the data loader in C++ and feeding that to Pytorch’s dataloader. That worked pretty well with respect to latency.</p>

Vinayaka_Hegde: <p><a class="mention" href="/u/lkc1">@lkc1</a> Thanks!<br>
Got it. I was able to reproduce the same with Jetson Orin 12 CPUs. Here are the results ;</p>
<p>Small obj case</p>
<pre><code class="lang-auto">avg time per batch for num workers=0: 0.459348852338735
avg time per batch for num workers=2: 0.2055776757770218
avg time per batch for num workers=4: 0.10303598642349243
avg time per batch for num workers=6: 0.07063477154588327
avg time per batch for num workers=8: 0.058643676224164665
</code></pre>
<p>Large obj case :</p>
<pre><code class="lang-auto">avg time per batch for num workers=0: 0.6697914574178867
avg time per batch for num workers=2: 0.44152931112330407
avg time per batch for num workers=4: 0.305355457123369
avg time per batch for num workers=6: 0.39800371177261695
avg time per batch for num workers=8: 0.3520681524532847
</code></pre>
<p>It would help me a lot to know where exactly the seralization/de-seralization happen in the mutiprocessing module in python. It would be great if you could point me to the lines of code. I would need to dig deeper into the exact overhead time caused by the seralization/de-serailization in the latter case when there are large objects. So, please let me know how can I achieve this? Thanks a lot !</p>

Arij-Aladel: <p><a class="mention" href="/u/vinayaka_hegde">@Vinayaka_Hegde</a></p>
<p>Didyou read the code?<br>
# prime the prefetch loop<br>
for _ in range(self._prefetch_factor * self._num_workers):<br>
self._try_put_index()<br>
then you have:</p>
<pre><code>def _try_put_index(self):
    assert self._tasks_outstanding &lt; self._prefetch_factor * self._num_workers

    try:
        index = self._next_index()
    except StopIteration:
        return
    for _ in range(self._num_workers):  # find the next active worker, if any
        worker_queue_idx = next(self._worker_queue_idx_cycle)
        if self._workers_status[worker_queue_idx]:
            break
    else:
        # not found (i.e., didn't break)
        return

    self._index_queues[worker_queue_idx].put((self._send_idx, index))
    self._task_info[self._send_idx] = (worker_queue_idx,)
    self._tasks_outstanding += 1
    self._send_idx += 1
</code></pre>
<p>Hope it is clear now</p>