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] = None,multiprocessing_context=None, generator=None,*, prefetch_factor: int = 2,persistent_workers: bool = False):

代码解析

在DataLoader的__init__函数里，我们可以看到，它实现了：

1. 判断是否是IterableDataset类型，如果是需要进一步判断参数是否正确
2. 构建Sampler，单样本
3. 构建BatchSampler，
4. 组建batch 构建collate
5. 其他的一些逻辑判断

IterableDataset 判断

• IterableDataset应用于数据集非常大，将其完全加载进内存不现实（例如高达几个TB的数据），这时就需要IterableDataset构建可迭代的Dataset类，自定义的Dataset需要继承自torch.util.data.IterableDataset，重写__iter__方法，返回可迭代对象（通常是yield生成器）
• 对于IterableDataset来说，就没有构建采样器Sampler的需求，因为样本是通过调用__iter__一个个读取出来的。执行封装的DataLoader传进去的batch_size次__iter__方法，就获取到一个mini batch

# 判断dataset是否是IterableDataset类型if isinstance(dataset, IterableDataset):self._dataset_kind = _DatasetKind.Iterable# 按照__iter__获取数据，所以不需要打乱if shuffle is not False:raise ValueError("DataLoader with IterableDataset: expected unspecified ""shuffle option, but got shuffle={}".format(shuffle))elif sampler is not None:# 按照__iter__获取数据，也不再需要sampler获取数据raise ValueError("DataLoader with IterableDataset: expected unspecified ""sampler option, but got sampler={}".format(sampler))elif batch_sampler is not None:# 按照__iter__获取数据，也不再需要batch_sampler获取数据索引raise ValueError("DataLoader with IterableDataset: expected unspecified " "batch_sampler option, but got batch_sampler={}".format(batch_sampler))else:self._dataset_kind = _DatasetKind.Map

构建Sampler，单样本

if sampler is None:  # give default samplersif self._dataset_kind == _DatasetKind.Iterable:# 如果是Iterable的Dataset,就采用迭代的方式获取samplersampler = _InfiniteConstantSampler()else:  # 否则判断是否使用shuffle，使用则随机产生sampler，不使用就按照顺序产生samplerif shuffle:sampler = RandomSampler(dataset, generator=generator)else:sampler = SequentialSampler(dataset)

构建BatchSampler，组建batch

• 注意，上面说batch_sampler不能和batch_size、sampler、drop_last同时使用是指：如果已经定义了batch_sampler则与batch_size和sampler互斥！！！前提是已经定义了batch_sampler！！！但是如果没有定义batch_sampler，则可以通过batch_size，sampler，dorp_last来组建batch！！！

# 要取batch_size个sampler，但是还没有取，即batch_sampler==None
if batch_size is not None and batch_sampler is None:# 获取batch_size个sampler个索引batch_sampler = BatchSampler(sampler, batch_size, drop_last)

构建collate_fn 对获取的batch进行处理

if collate_fn is None:if self._auto_collation:# 默认的实际上什么也没干collate_fn = _utils.collate.default_collateelse:collate_fn = _utils.collate.default_convert

其他的一些逻辑判断

# sampler 不能和 shuffle 同时出现
# 因为shuffle是将数据打乱，而sampler是按照某一顺序获取数据if sampler is not None and shuffle:raise ValueError('sampler option is mutually exclusive with ''shuffle')if batch_sampler is not None:# batch_sampler不能和batch_size,shuffle,sampler,drop_last同时使用。# batch_sampler可以理解为batch_size和sampler的结合if batch_size != 1 or shuffle or sampler is not None or drop_last:raise ValueError('batch_sampler option is mutually exclusive ''with batch_size, shuffle, sampler, and '                          'drop_last')batch_size = Nonedrop_last = Falseelif batch_size is None:# batch_size为None，默认是1，如果drop_last为True就会舍弃最后一个，这样数据就会减少。（构成了一个batch但是仍然舍弃掉）if drop_last:raise ValueError('batch_size=None option disables auto-batching ''and is mutually exclusive with drop_last')self.collate_fn = collate_fnself.persistent_workers = persistent_workersself.__initialized = Trueself._IterableDataset_len_called = None  # See NOTE [ IterableDataset and __len__ ]self._iterator = Noneself.check_worker_number_rationality()torch.set_vital('Dataloader', 'enabled', 'True')  # type: ignore[attr-defined]

_get_iterator

代码解析

def _get_iterator(self) -> '_BaseDataLoaderIter':if self.num_workers == 0:# 单线程return _SingleProcessDataLoaderIter(self)else:# 多线程self.check_worker_number_rationality()return _MultiProcessingDataLoaderIter(self)

multiprocessing_context

multiprocessing_context

setattr

iter

代码解释

 # 其中 -> '_BaseDataLoaderIter' 是函数注释，运行时跟没有加注解之前的效果也没有任何差距。# 主要作用是提醒程序猿这里应该是 '_BaseDataLoaderIter'的数据类型def __iter__(self) -> '_BaseDataLoaderIter':if self.persistent_workers and self.num_workers > 0:if self._iterator is None:self._iterator = self._get_iterator()else:self._iterator._reset(self)return self._iteratorelse:return self._get_iterator()

_auto_collation

代码解析

	@propertydef _auto_collation(self):# 根据batch_sampler判断是否设置_auto_collationreturn self.batch_sampler is not None

_index_sampler

len

check_worker_number_rationality

_SingleProcessDataLoaderIter

代码解析

def __init__(self, loader):super(_SingleProcessDataLoaderIter, self).__init__(loader)assert self._timeout == 0assert self._num_workers == 0self._dataset_fetcher = _DatasetKind.create_fetcher(self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last)def _next_data(self):# 获取索引index = self._next_index()  # may raise StopIteration# 获取数据data = self._dataset_fetcher.fetch(index)  # may raise StopIterationif self._pin_memory:data = _utils.pin_memory.pin_memory(data)# 返回数据return data

_BaseDataLoaderIter

__next__方法会调用_next_data，_next_data获取一个batch的数据