Class DataLoader

imperative/python/megengine/data/dataloader.py:89–245

r"""Data loader. Combines a dataset and a sampler, and provides a convenient way to iterate on a given dataset. The process is as follows: .. mermaid:: :align: center flowchart LR Dataset.__len__ -- Sampler --> Indices batch_size -- Sampler --> Indices

class DataLoader:
    r"""Data loader. Combines a dataset and a sampler, and provides a convenient way
    to iterate on a given dataset. The process is as follows:

    .. mermaid::
       :align: center
    
       flowchart LR
          Dataset.__len__ -- Sampler --> Indices
          batch_size -- Sampler --> Indices
          Indices -- Dataset.__getitem__ --> Samples
          Samples -- Transform + Collator --> mini-batch

    See :ref:`data-guide` for more details.

    Args:
        dataset: dataset from which to load the minibatch.
        sampler: defines the strategy to sample data from the dataset.
            If ``None``, it will sequentially sample from the dataset one by one.
        transform: defined the transforming strategy for a sampled batch.
        collator: defined the merging strategy for a transformed batch.
        num_workers: the number of sub-process to load, transform and collate
            the batch. ``0`` means using single-process. Default: 0
        timeout: if positive, means the timeout value(second) for collecting a
            batch from workers. Default: 0
        preload: whether to enable the preloading strategy of the dataloader. 
            When enabling, the dataloader will preload one batch to the device memory to speed up the whole training process.
        parallel_stream: whether to splitting workload across all workers when dataset is streamdataset and num_workers > 0.
            When enabling, each worker will collect data from different dataset in order to speed up the whole loading process.
            See ref:`streamdataset-example` for more details

    Examples:
        >>> import megengine.data as data
        >>> dataset = CustomDataset()         # doctest: +SKIP
        >>> dataloader = DataLoader(dataset)  # doctest: +SKIP
        >>> for batch_data in DataLoader:     # doctest: +SKIP
        >>>     print(batch_data.shape)       # doctest: +SKIP

    .. admonition:: The effect of enabling preload
       :class: warning

       * All elements in :class:`map`, :class:`list`, and :class:`tuple` will be converted to :class:`~.Tensor` by preloading,
         and you will get :class:`~.Tensor` instead of the original Numpy array or Python built-in data structrure.
       * Tensors' host2device copy and device kernel execution will be overlapped,
         which will improve the training speed at the cost of **higher device memory usage** (due to one more batch data on device memory).
         This feature saves more time when your NN training time is short or your machine's host PCIe bandwidth for each device is low.
    """

    def __init__(
        self,
        dataset: Dataset,
        sampler: Sampler = None,
        transform: Transform = None,
        collator: Collator = None,
        num_workers: int = 0,
        timeout: int = 0,
        preload: bool = False,
        parallel_stream: bool = False,