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)       

启用预加载功能的后果

通过预加载功能 map，list`和 :class:`tuple`都将会被转为 :class:`~.Tensor ，并且最后获得也是 :class:`~.Tensor`而不是numpy数组或python内置数据结构
张量的host2device的过程会合设备kernel的执行过程并行起来，以提高训练的速度，代价是**更高的设备内存使用**(由于设备内存上多了一个批处理数据)。当神经网络训练时间较短或机器上每个设备的PCIe带宽较低时，该特性可以节省更多的时间

megengine.data

Dataset