DataLoader

class DataLoader(dataset, sampler=None, transform=None, collator=None, num_workers=0, timeout=0, preload=False, parallel_stream=False)[源代码]

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

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

See 使用 Data 构建输入 Pipeline for more details.

参数
  • dataset (Dataset) – 需要从中分批加载的数据集。

  • sampler (Optional[Sampler]) – 定义从数据集上采样数据的策略。如果设为 None,将从数据集中逐个采样

  • transform (Optional[Transform]) – 为采样数据定义转换策略

  • collator (Optional[Collator]) – 为处理后的batch数据定义合并策略

  • num_workers (int) – 加载、转换和整理批次的子进程数量。 0 表示使用单进程。 默认:0

  • timeout (int) – 如果为正,则表示从 worker 那里收集批次的超时值(秒)。 默认:0

  • preload (bool) – 是否启用dataloader的预加载策略。如果启用时,dataloader将预加载一个batch的数据到设备上,以加快整个训练过程

  • parallel_stream (bool) – 当数据集是流数据集(streamdataset)并且 num_workers > 0时,是否将工作负载分配给所有workers。启用时,每个worker将从不同的数据集收集数据,以加快整个加载流程。更多细节请看:streamdataset-example

实际案例

>>> import megengine.data as data
>>> dataset = CustomDataset()         
>>> dataloader = DataLoader(dataset)  
>>> for batch_data in DataLoader:     
>>>     print(batch_data.shape)       

启用预加载功能的后果

  • 通过预加载功能 maplist`和 :class:`tuple`都将会被转为 :class:`~.Tensor ,并且最后获得也是 :class:`~.Tensor`而不是numpy数组或python内置数据结构

  • 张量的host2device的过程会合设备kernel的执行过程并行起来,以提高训练的速度,代价是**更高的设备内存使用**(由于设备内存上多了一个批处理数据)。当神经网络训练时间较短或机器上每个设备的PCIe带宽较低时,该特性可以节省更多的时间