以上就是最抽象的 DataLoader 运行流程以及和 Dataset、Sampler、BatchSampler、collate_fn 的关系。

首先需要强调的是 Dataset、Sampler、BatchSampler 和 DataLoader 都直接或间接实现了迭代器。

Dataset通过__getitem__方法使其可迭代

Sample对象是一个可迭代的基类对象，其常用子类 SequentialSampler 在 __iter__ 内部返回迭代器，RandomSampler 在 __iter__ 内部通过 yield 关键字返回迭代器

Batchsampler也是在__iter__内部通过yield关键字返回迭代器

DataLoader通过__iter__和__next__直接实现迭代器

除了DataLoader本身是迭代器外，其余对象本身都不是迭代器，但可以for in迭代

由于 DataLoader 类写的非常通用，故 Dataset、Sampler、BatchSampler 都可以外部传入，除了 Dataset 必须输入外，其余两个类都有默认实现，最典型的 Sampler 就是 SequentialSampler 和 RandomSampler。

需要注意的是 Sampler 对象其实在大部分时候都不需要传入 Dataset 实例对象，因为其功能仅仅是返回索引而已，并没有直接接触数据。

三、整体框架的讲解

核心运行逻辑：

def __next__(self):#返回batch个索引index=next(self.batch_sampler)#利用索引去取数据data=[self.dataset[idx] for idx in index]#batch维度聚合data=self.collate_fn(data)return data

整体流程：

1.self.batch_sampler=iter(batch_sampler)。在DataLoader的类初始化，需要得到BatchSampler的迭代器对象。

2.index=next(self.batch_sampler)。对于每次迭代，DataLoader对象首先会调用BatchSampler的迭代器进行下一次迭代，具体是调用BatchSampler对象的__iter__方法

3.而BatchSampler对象的__iter__方法实际上是需要依靠Sampler对象进行迭代输出索引，Sampler对象也是一个迭代器，当迭代batch_size次后就可以得到batch_size个数据索引。

4.data=[self.dataset[idx] for idx in index]。有了batch个索引就可以通过不断调用dataset的__getitem__方法返回数据对象，此时data就包含了batch个对象。

5.data=self.collate_fn(data)。将batch个对象输入给聚合函数，在第0个维度也就是batch维度进行聚合，得到类似（batch，....）的对象。

6.重复上面的操作，就可以不断输出一个一个的batch数据

class Dataset(object):#只要实现了__getitem__方法就可以变成迭代器def __getitem__(self,index):raise NotImplementedErrordef __len__(self):raise NotImplementedError

class Sampler(object):def __init__(self,data_source):passdef __iter__(self):raise NotImplementedErrordef __len__(self):raise NotImplementedError
#一般出现raise NotImplementedError这个错误，就是子类没有重写父类中的成员函数，然后子类对象调用此函数会报这个错误class SequentialSampler(sampler):def __init__(self,data_source):super(SequentialSampler,self).__init__(data_source)self.data_source=data_sourcedef __iter__(self):#返回迭代器，不然无法for  inreturn iter(range(len(self.data_source))def __len__(self):return len(self.data_source)class BatchSampler(Sampler):def __init__(self,sampler,batch_size,drop_last):self.sampler=samplerself.batch_size=batch_sizeself.dorp_last=drop_lastdef __iter__(self):batch=[]for idx in self.sampler:batch.append(idx)#如果得到了batch个索引，则可以通过yield关键字生成生成器返回，得到迭代器对象if len(batch)==self.batch_size:yield batchbatch=[]if len(batch)>0 and not self.drop_last:yield batchdef __len__(self):if self.drop_last:#如果最后的索引数不等于一个batch，抛弃return len(self.sampler)//self.batch_sizeelse:return (len(self.sampler)+self.batch_size-1)//self.batch_size

class DataLoader(object):def __init__(self,dataset,batch_size=1,shuffle=False,sample=None,batch_sampler=None,collate_fn=None,drop_last=False):self.dataset=dataset#因为这两个功能是冲突的if sampler is not None and shuffle:raise ValueError('sampler option is ..')if batch_sampler is not None:# 一旦设置了 batch_sampler，那么 batch_size、shuffle、sampler# 和 drop_last 四个参数就不能传入# 因为这4个参数功能和 batch_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 = Falseif sampler is None:if shuffle:sampler = RandomSampler(dataset)else:sampler = SequentialSampler(dataset)# 也就是说 batch_sampler 必须要存在，你如果没有设置，那么采用默认类if batch_sampler is None:batch_sampler = BatchSampler(sampler, batch_size, drop_last)self.batch_size = batch_sizeself.drop_last = drop_lastself.sampler = samplerself.batch_sampler = iter(batch_sampler)if collate_fn is None:collate_fn = default_collateself.collate_fn = collate_fn#核心代码def __next__(self):index=next(self.batch_sampler)data=[self.dataset[idx] for idx in index]data=self.collate_fn(data)return data#返回自身，因为自身实现了nextdef __iter__(self):return self

def default_collate(batch):elem=batch[0]elem_type=type(elem)if isinstance(elem,torch.Tensor):return torch.stack(batch,0)elif elem_type.__module__=='numpy':return default_collate([torch.as_tensor(b) for b in batch])else:raise NotImplementedError

完整调用例子

class Simplev1Dataset(Dataset):def __init__(self):#伪造数据self.imgs=np.arange(0,16).reshape(8,2)def __getitem__(self,index):return self.imgs[index]def __len__(self):return self.imgs.shape[0]from simplev1_dataset import Simplev1Dataset
simple_dataset=Simplev1Dataset()
dataloader=DataLoader(simple_dataset,batch_size=2,collate_fn=default_collate)
for data in dataloader:print(data)

四、Reference

https://zhuanlan.zhihu.com/p/340465632