Unexpected parallel data loader performance using IterableDatasets compared to (map-style) Datasets with num_workers > 1

username: Christos_Josephides
cooked: <p>Hi,</p>
<p>I’m trying to diagnose a performance discrepancy between using IterableDatasets and (map-style)Datasets in a multi-processed data loader setting.</p>
<p>My experiment (code at the end of the post) consisted of:</p>
<ol>
<li>Make a (Iterable)Dataset. This synthesizes dummy data and possibly adds a time delay to simulate batch loading work.</li>
<li>Make a Dataloader that consumes the above dataset instance. The number of workers varied from 0 (loading in the main process) to 4.</li>
<li>Iterate through 100 batches of data yielded by the dataloader, with possibly a time delay to simulate batch processing work (e.g. training).</li>
</ol>
<p>I varied the number of workers, batch loading time, and batch processing time. I plotted my results as</p>
<h2><a name="iterabledataset-left-vs-map-styledataset-right-1" class="anchor" href="#iterabledataset-left-vs-map-styledataset-right-1"></a>IterableDataset (left) vs (map-style)Dataset (right)</h2>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b.png" data-download-href="https://discuss.pytorch.org/uploads/default/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b" title="dataset_comparison"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b_2_636x500.png" alt="dataset_comparison" data-base62-sha1="y9ZTlxier2JhMtrUDFo8WbOWeVB" width="636" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b_2_636x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b_2_954x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef6b3f6d14bad64a29e77aa95e97e0b945c0369b_2_1272x1000.png 2x" data-dominant-color="FBFCFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">dataset_comparison</span><span class="informations">2544×2000 228 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The vertical axis is time to run through N=100 batches divided by the time to do so with num_workers = 0 in each case. I grouped the results by batch processing / loading time.</p>
<p>I expected the behavior of (map-style)Datasets, but I did not expect a linear degradation in performance for num_workers &gt;= 1 when using IterableDatasets.</p>
<p>Could someone please have a look at my code below and let me know if I’m using IterableDataset incorrectly or if there might be a bug we should know about/report?</p>
<p>Thanks,<br>
Christos</p>
<pre><code class="lang-auto">import os
import time
import itertools
import numpy as np
import torch


class TimedBlock:
    """
    Context manager to measure wall time.
    """

    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.t_start = time.perf_counter()

    def __exit__(self, *args):
        print("(%d) %s execution finished. Time elapsed = %.3fs." % (os.getpid(), self.name,
                                                                     time.perf_counter() - self.t_start))


class IterableDataset(torch.utils.data.IterableDataset):
    """
    Synthetic iterable dataset that simulates data loading work.
    """

    def __init__(self, n_items, shape, batch_loading_time):
        """
        Args:
            n_items: number of items to return.
            shape: shape of batch tensor to return.
            batch_loading_time: each batch will take at least this many seconds to return.
        """
        self.n_items = n_items
        self.shape = shape
        self.batch_loading_time = batch_loading_time

    def __iter__(self):
        self.i = 0
        while self.i &lt; self.n_items:
            t_start = time.perf_counter()
            X = np.random.randn(*self.shape)
            t_remaining =  self.batch_loading_time - (time.perf_counter() - t_start)
            if t_remaining &gt; 0:
                time.sleep(t_remaining)
            yield X
            self.i += 1


class Dataset(torch.utils.data.Dataset):
    """
    Synthetic map-style dataset that simulates data loading work.
    """

    def __init__(self, n_items, shape, batch_loading_time):
        """
        Args:
            n_items: number of items to return.
            shape: shape of batch tensor to return.
            batch_loading_time: each batch will take at least this many seconds to return.
        """
        self.n_items = n_items
        self.shape = shape
        self.batch_loading_time = batch_loading_time

    def __len__(self):
        return self.n_items

    def __getitem__(self, idx):
        t_start = time.perf_counter()
        X = np.random.randn(*self.shape)
        t_remaining =  self.batch_loading_time - (time.perf_counter() - t_start)
        if t_remaining &gt; 0:
            time.sleep(t_remaining)
        return X


def test_simpleloader(n_iters=int(1e2), shape=(1,), num_workers=0, batch_loading_time=0, batch_process_time=0,
                      exp_name="simpleloader", dataset_cls=IterableDataset):
    """
    Load and process n_iters batches of data using a dataloader and one of the above dataset classes. We simulate
    training by waiting batch_process_time seconds per batch.
    """

    dataset = dataset_cls(n_iters, shape, batch_loading_time)
    data_loader = torch.utils.data.DataLoader(dataset,
                                             batch_size=None,
                                             num_workers=num_workers,
                                             collate_fn=None,
                                             pin_memory=False,
                                             worker_init_fn=None,
                                             multiprocessing_context=None)
    
    with TimedBlock(exp_name):
        for i, x in enumerate(data_loader):
            time.sleep(batch_process_time)


def simpleloader_grid(num_workers_grid=[0, 1, 2, 3, 4],
                      batch_loading_time_grid=[10**i for i in range(0, -2, -1)],
                      batch_process_time_grid=[10**i for i in range(0, -2, -1)],
                      dataset_cls=IterableDataset):

    n_iters = int(1e2)
    shape = (1,)

    params = itertools.product(num_workers_grid, batch_loading_time_grid, batch_process_time_grid)

    for p in params:
        test_simpleloader(n_iters, shape, p[0], p[1], p[2], str(p), dataset_cls=dataset_cls)


if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn")  # probably not important here
    t_start = time.perf_counter()
    simpleloader_grid(dataset_cls=IterableDataset)
    print("Main execution time = %.3fs" % (time.perf_counter() - t_start))

</code></pre>