【代码笔记】Pytorch学习 DataLoader模块详解

dataloader整体结构

dataloader主要有6个class构成（可见下图）

• _DatasetKind：
• _InfiniteConstantSampler：
• DataLoader：
• _BaseDataLoaderIter：
• _SingleProcessDataLoaderIter：
• _MultiProcessingDataLoaderIter：
  

DataLoader

我们首先看一下DataLoader的整体结构：

• init:
• _get_iterator:
• multiprocessing_context:
• multiprocessing_context:
• setattr:
• iter:
• _auto_collation:
• _index_sampler:
• len:
• check_worker_number_rationality:

init 初始化

参数解释

这里会把参数全部列出，这里列出的目的是让大家知道各个参数的意义。实际上很多是用不到的，我用加粗字体表示一些常用的参数。

• self：代之Dataset这个类本身
• dataset: Dataset[T_co]是默认值，是你要处理的数据集
• batch_size: Optional[int] = 1, 可选，默认是1。每个batch可以加载batct_size个数据。
• shuffle: bool = False, 每轮训练后是否将数据集打乱
• sampler: Optional[Sampler] = None, 默认是None 自定义方法（某种顺序）从Dataset中取样本，指定这个参数就不能设置shuffle。因为shuffle是打乱数据集的顺序，而sample是以某种顺序取数据，所以二者互斥！sampler可能是获取一整个数据集的数据，是对一整个数据集进行操作，而不是一个batch_size。
• batch_sampler: Optional[Sampler[Sequence]] = None, 返回一个batch的索引，与batch_size, shuffle, sampler, drop_last互斥
  传入了batch_sampler，相当于已经告诉了PyTorch如何从Dataset取多少数据，怎么取数据去组成一个mini batch，所以不需要以上参数。可以理解为batch_sampler是batch_size和sampler的结合，所以不需要batch_size, sampler, shuffle, drop_last（因为drop_last也是怎么取数据）。
• num_workers: int = 0, 多进程加载数据，默认为0，即采用主进程加载数据
• collate_fn: Optional[_collate_fn_t] = None, 聚集函数，用来对一个batch进行后处理，拿到一个batch的数据后进行什么处理，返回处理后的batch数据。默认源码中进行了若干逻辑判断，仅将数据组合起来返回，没有实质性工作。默认collate_fn的声明是：def default_collate(batch): 所以自定义collate_fn需要以batch为输入，以处理后的batch为输出。类似于transform,transform是对单个数据处理，而collate_fn是对单个batch做处理。
• pin_memory: bool = False, 用于将tensor加载到GPU中进行运算
• drop_last: bool = False, 是否保存最后一个mini batch，样本数量可能不支持被batch size整除，所以drop_last参数决定是否保留最后一个可能批量较小的batch
• timeout: float = 0, 控制从进程中获取一个batch数据的时延
• worker_init_fn: Optional[_worker_init_fn_t] = None, 初始化子进程
• multiprocessing_context=None,
• generator=None,
• prefetch_factor: int = 2, 控制样本在每个进程里的预加载，默认为2
• persistent_workers: bool = False 控制加载完一次Dataset是否保留进程，默认为False

def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1,
                 shuffle: bool = False, sampler: Optional[Sampler] = None,
                 batch_sampler: Optional[Sampler[Sequence]] = None,
                 num_workers: int = 0, collate_fn: Optional[_collate_fn_t] = None,
                 pin_memory: bool = False, drop_last: bool = False,
                 timeout: float = 0, worker_init_fn: Optional[_worker_init_fn_t]