Class ThreadDataLoader

monai/data/thread_buffer.py:137–199

Subclass of `DataLoader` using a `ThreadBuffer` object to implement `__iter__` method asynchronously. This will iterate over data from the loader as expected however the data is generated on a separate thread. Use this class where a `DataLoader` instance is required and not just an iter

class ThreadDataLoader(DataLoader):
    """
    Subclass of `DataLoader` using a `ThreadBuffer` object to implement `__iter__` method asynchronously. This will
    iterate over data from the loader as expected however the data is generated on a separate thread. Use this class
    where a `DataLoader` instance is required and not just an iterable object.

    The default behaviour with `repeats` set to 1 is to yield each batch as it is generated, however with a higher
    value the generated batch is yielded that many times while underlying dataset asynchronously generates the next.
    Typically not all relevant information is learned from a batch in a single iteration so training multiple times
    on the same batch will still produce good training with minimal short-term overfitting while allowing a slow batch
    generation process more time to produce a result. This duplication is done by simply yielding the same object many
    times and not by regenerating the data.

    Another typical usage is to accelerate light-weight preprocessing (usually cached all the deterministic transforms
    and no IO operations), because it leverages the separate thread to execute preprocessing to avoid unnecessary IPC
    between multiple workers of DataLoader. And as CUDA may not work well with the multi-processing of DataLoader,
    `ThreadDataLoader` can be useful for GPU transforms. For more details:
    https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md.

    The `use_thread_workers` will cause workers to be created as threads rather than processes although everything else
    in terms of how the class works is unchanged. This allows multiple workers to be used in Windows for example, or in
    any other situation where thread semantics is desired. Please note that some MONAI components like several datasets
    and random transforms are not thread-safe and can't work as expected with `thread workers`, need to check all the
    preprocessing components carefully before enabling `use_thread_workers`.

    See:
        * Fischetti et al. "Faster SGD training by minibatch persistency." ArXiv (2018) https://arxiv.org/abs/1806.07353
        * Dami et al., "Faster Neural Network Training with Data Echoing" ArXiv (2020) https://arxiv.org/abs/1907.05550
        * Ramezani et al. "GCN meets GPU: Decoupling "When to Sample" from "How to Sample"." NeurIPS (2020).
          https://proceedings.neurips.cc/paper/2020/file/d714d2c5a796d5814c565d78dd16188d-Paper.pdf

    Args:
        dataset: input dataset.
        buffer_size: number of items to buffer from the data source.
        buffer_timeout: time to wait for an item from the buffer, or to wait while the buffer is full when adding items.
        repeats: number of times to yield the same batch.
        use_thread_workers: if True and num_workers > 0 the workers are created as threads instead of processes
        kwargs: other arguments for `DataLoader` except for `dataset`.

    """

    def __init__(
        self,
        dataset: Dataset,
        buffer_size: int = 1,
        buffer_timeout: float = 0.01,
        repeats: int = 1,
        use_thread_workers: bool = False,
        **kwargs,
    ):
        # if workers should be threads, create a new multiprocessing context with the process and queue types
        # substituted with the shim types given above
        if use_thread_workers and kwargs.get("num_workers", 0) > 0:
            kwargs["multiprocessing_context"] = _ProcessThreadContext()
            kwargs["persistent_workers"] = False

        super().__init__(dataset, **kwargs)
        self.buffer_size = buffer_size