不理解 看代码 看代码 看代码 重要的事情说三遍 代码最好进行调试理解 代码是最好的老师

1 迭代器的简要介绍

1.1 iter和next的初步理解

__iter__(): class 中定义的魔术方法
__next__(): class 中定义的魔术方法
iter(): 将一个定义了__iter__()的类转化为一个迭代器
next(): 读取迭代器中的数据

iter()和next()的原理图可以参考这个链接：Pytorch（三）：Dataset和Dataloader的理解

注意：当一个类中同时定义了__iter__()和__next__()，这个类本身就是一个迭代器，或者如果只定义了__iter__()，需要使用iter()将类对象转化为迭代器

next()有一个点需要注意(见下面代码)：

1. 类Exmaple同时定义了__iter__()和__next__() => exmaple是一个迭代器，next(example)调用的是example中的__next__()
2. 类Exmaple只定义了__iter__() => 需要用iter(example)将exmaple转化为一个迭代器，next(example)调用的是example中的__iter__()

class Example:def __init__(self, data):self.data = datadef __iter__(self):for index in range(len(self.data)):print("打印iter")yield self.data[index]def __next__(self):print("打印next")if __name__ == "__main__":example = Example(range(100))"""不注释__next__()  Exmaple同时定义了__iter__()和__next__() => exmaple是一个迭代器当Example同时定义了__iter__()和__next__()的时候 => next(example)调用的是example中的__next__()"""next(example)  """注释__next__()  Exmaple只定义了__iter__() => 需要用iter(example)将exmaple转化为一个迭代器当Example 只定义了 __iter__()的时候 => next(example)调用的是example中的__iter__()"""# next(iter(example))  

for循环：当执行for i in obj:，Python解释器会在第一次迭代时自动调用iter(obj)生成一个迭代器Iter，然后在之后的迭代(包扩第一次迭代)会使用next()调用迭代器Iter的__iter__()方法，for语句会自动处理最后抛出的StopIteration异常，如list、tuple、dict等可迭代对象。但是当obj是一个类对象，且类中定义了__iter__()时，Python解释器会在循环开始时自动调用Iter = obj.__iter__()返回一个迭代器Iter，然后在之后的迭代(包括第一次迭代)会调用迭代器Iter的__next__()方法，for语句会自动处理最后抛出的StopIteration异常。

如果obj是类似于list、tuple、dict的一个可迭代对象时，python解释器会自动为我们做以下几步操作(通义千问，重点不是for循环中这个用法，不用太过重视)：

• 调用 iter() 函数：首先，Python调用内置的iter()函数尝试从可迭代对象obj中获取一个迭代器。这相当于执行了iterator = iter(obj)。
• 循环与 next() 方法：在for循环的每次迭代中，Python会使用next()调用迭代器的__iter__()方法来获取下一个元素。
• 处理 StopIteration 异常：当迭代器没有更多的元素可以提供时，next()方法会抛出一个StopIteration异常。Python的for循环会捕获这个异常，并且知道这是迭代结束的信号，于是它会干净利落地结束循环，而不会将这个异常暴露给用户代码。
• 具体的实现代码过程类似于

  for i in obj:# 循环体
  

  实际上，Python在幕后做的事情类似于下面的手动迭代过程：

  iterator = iter(obj)
  while True:try:i = next(iterator)except StopIteration:break  # 遇到StopIteration时退出循环# 循环体
  

1.2 torch.utils.data.DataLoader中的iter和next

DataLoader(torch.utils.data.DataLoader)重要的是：在数据集遍历的一般化过程中，执行for i, data in enumerate(dataLoader):，当dataLoader是一个类对象，且类中定义了__iter__()时，Python解释器会在循环开始时自动调用Iter = dataLoader.__iter__()返回一个迭代器Iter，然后在之后的迭代(包扩第一次迭代)会调用迭代器Iter的__next__()方法获取数据。

__iter__()的作用是使一个类对象成为可迭代的。在for循环中，Python调用iter(obj)时，如果类对象obj是可迭代的（即定义了__iter__），那么iter(obj)实际上等同于调用obj.__iter__()生成一个迭代器。 这个方法返回的迭代器随后被用于每次循环迭代中调用__next__()来获取下一个值，直到遇到StopIteration异常为止。

1.2.1 第一种方式

下面代码中，DataLoader的load_data方法通过yield from委托给了iter(self.sampler)的这个迭代器，从而直接迭代并输出数据。注意yeild from iter(self.sampler)中必须要使用iter()，因为self.sampler并不是一个迭代器(Sampler中只定义了__iter__())，需要使用iter(self.sampler)将self.sampler转换为迭代器。

# 打印0-99
class Sampler:def __init__(self, data):self.data = datadef __iter__(self):for index in range(len(self.data)):yield self.data[index]class DataLoader:def __init__(self, sampler):self.sampler = samplerdef load_data(self):yield from iter(self.sampler)# 使用示例
sampler = Sampler(range(100))  # 假设的采样器实例
data_loader = DataLoader(sampler)
for item in data_loader.load_data():print(item)

1.2.2 第二种方式，这种方式就类似于torch.utils.data.DataLoader。

1. 在for item in data_loader:中，会调用data_loader.__iter__()返回一个迭代器Iter，然后一直调用迭代器Iter的__next__()返回值。在for循环中的第一次循环调用data_loader.__iter__()生成一个迭代器Iter后，之后的循环就不再调用data_loader.__iter__()了，只调用迭代器Iter的__next__()读取数据。当然，如果再次使用for循环，它又会再次调用data_loader.__iter__()生成一个迭代器。如目标检测中，使用torch.utils.data.DataLoader读取数据，epochs=300时，每一个epoch，都会使用一次for循环，即都会调用一次__iter__()生成一个迭代器。
2. 因为class DataLoader同时定义了__iter__()和__next__()，所以DataLoader类本身就是一个迭代器，所以在其__iter__()中可以使用return self返回其本身这个迭代器Iter， 然后调用迭代器Iter的Iter.__next__()去读取数据。注意，这里不要理解错了，调用的__next__()是迭代器Iter的，如果class DataLoader中的__iter__()返回的不是self，而是其他迭代器如iter(self.sampler)，那么就不会调用class DataLoader的__next__()，而是调用iter(self.sampler)的__next__()。
3. 还有一个注意的地方，这个在torch.utils.data.DataLoader中也出现过，就是使用next(self.sampler)时一定要注意self.sampler必须是一个迭代器，因为sampler不是一个迭代器，所以需要在__init__()使用iter(sampler)将sampler转化成迭代器。

   # 打印0-99
   class Sampler:def __init__(self, data):self.data = datadef __iter__(self):for index in range(len(self.data)):yield self.data[index]class DataLoader:def __init__(self, sampler):self.sampler = iter(sampler)def __iter__(self):return selfdef __next__(self):return next(self.sampler)# 使用示例
   sampler = Sampler(range(100))  # 假设的采样器实例
   data_loader = DataLoader(sampler)
   for item in data_loader:print(item)
   

1.2.3 第三种方式，这种方式就是torch.utils.data.DataLoader的实现方式

1. 在for item in data_loader:中，会调用data_loader.__iter__()返回一个迭代器Iter，这个迭代器Iter就是DataLoaderIter(self)。self是DataLoader作为参数传入到DataLoaderIter进行初始化。注意：DataLoaderIter(self)中的__iter__()只是为了让DataLoaderIter成为一个迭代器，实际运行过程中并不调用这个__iter__()。
2. 在生成迭代器DataLoaderIter之后，会调用迭代器DataLoaderIter中的__next__()。

   # 打印0-99
   # Sampler 类似于 BatchSampler
   class Sampler:def __init__(self, data):self.data = datadef __iter__(self):for index in range(len(self.data)):yield self.data[index]# DataLoaderIter 类似于 _BaseDataLoaderIter
   class DataLoaderIter:def __init__(self, loader):self.sampler = iter(loader.sampler)def __iter__(self):return selfdef __next__(self):return next(self.sampler)# DataLoader 类似于 DataLoader
   class DataLoader:def __init__(self, sampler):self.sampler = samplerdef __iter__(self):return DataLoaderIter(self)# 使用示例
   if __name__ == "__main__":sampler = Sampler(range(100))  # 假设的采样器实例data_loader = DataLoader(sampler)epochs=3for i in range(epochs):for item in data_loader:print(item)
   

1.2.4 第四种方式，这种方式就是YOLOv9中InfiniteDataLoader的实现方式

运行流程：data_loader.__iter__() => InfiniteDataLoader.__iter__() => next(self.iterator) => DataLoaderIter.__next__() => next(self.sampler) => _RepeatSampler.__iter__() => Sampler.__iter__()

1. 首先data_loader = InfiniteDataLoader()：在InfiniteDataLoader中，语句object.setattr(self, ‘sampler’, _RepeatSampler(self.sampler))含义是先执行self.sampler = Sampler(range(100))， 然后执行self.sampler = _RepeatSampler(self.sampler)。self.iterator = super().__iter__()的含义是self.iterator = DataLoaderIter(self)。
2. 然后对于for item in data_loader:，首先调用data_loader.__iter__()生成一个迭代器。for _ in range(len(self)):中，len(self) => len(self.sampler.sampler) => len(self.data)。yield next(self.iterator)中yield就是生成一个迭代器，由于DataLoaderIter中定义了__next__()，next(self.iterator)调用的就是DataLoaderIter.__next__()得到next(self.sampler)。
3. 在步骤1中已经知道self.sampler = _RepeatSampler(self.sampler)，next(self.sampler)此时调用的是_RepeatSampler.__iter__()。关于_RepeatSampler.__iter__()中while True无限循环读取数据可以查看1.2.4.1小节
4. yeild from 委托给Sampler.__iter__()，因此Sampler.__iter__()控制流程产出元素。

章节1.2.4使用self.iterator好处：在(目标检测中的epoch)每一次epoch时，都不需要重新进行初始化迭代器DataLoaderIter(self)；而对于章节1.2.3中的每一epoch，for循环都会调用data_loader.__iter__()，返回的是DataLoaderIter(self)，因此每一个epoch迭代器DataLoaderIter(self)都会重新初始化。

# 打印三遍0-99
class Sampler:def __init__(self, data):self.data = datadef __len__(self):return len(self.data)def __iter__(self):for index in range(len(self.data)):yield self.data[index]class DataLoaderIter:def __init__(self, loader):# self.sampler = iter(loader.sampler)self.sampler = iter(loader.sampler)def __iter__(self):return selfdef __next__(self):return next(self.sampler)class DataLoader:def __init__(self):self.sampler =  Sampler(range(100))  # 假设的采样器实例def __iter__(self):return DataLoaderIter(self)class InfiniteDataLoader(DataLoader):""" Dataloader that reuses workersUses same syntax as vanilla DataLoader"""def __init__(self, *args, **kwargs):super().__init__(*args, **kwargs)object.__setattr__(self, 'sampler', _RepeatSampler(self.sampler))self.iterator = super().__iter__()def __len__(self):return len(self.sampler.sampler)def __iter__(self):for _ in range(len(self)):yield next(self.iterator)class _RepeatSampler:""" Sampler that repeats foreverArgs:sampler (Sampler)"""def __init__(self, sampler):self.sampler = samplerdef __iter__(self):while True:yield from iter(self.sampler)# 使用示例
if __name__ == "__main__":data_loader = InfiniteDataLoader()epochs=3for i in range(epochs):for item in data_loader:print(item)

1.2.4.1 while True

这个是章节1.2.4代码的简化版，是为了理解，为什么注释掉while True后，外层循环遍历一个epoch后(只会打印一遍0-99)，会出现报错。

这个问题的理解关键其实是yeild from：
yeild from 用于在一个生成器中委托给另一个可迭代对象进行迭代。换句话说，它使得当前生成器暂停自己的执行，并允许另一个迭代器控制流程，直接产出其元素。当那个迭代器完成（抛出 StopIteration 异常）时，yield from 表达式的执行结束，控制权返回到调用方。

代码理解，就是next(self.loader)在调用DataLoaderIte.__iter__()迭代器时，将其委托给Sampler.__iter__()，Sampler.__iter__()直接产出元素，当Sampler.__iter__()迭代器完成后，yeild from yield from 表达式的执行结束，控制权返回DataLoaderIte.__iter__()。

可以在for i in range(epcohs):处打一个断点进行调试，第一个epoch会顺利打印0-99，在第二个epoch开始，首先会跳到语句yield next(self.loader)，此时进行单步调试后会直接跳到for index in range(len(self.data)):，跳过了yield from iter(self.sampler)，index的值是99，也就代表Sampler.__iter__()迭代器还没有迭代完成。由于index=99, self.data=100，再次进行单步调试后，Sampler.__iter__()迭代器就迭代完成了，会跳到yield from iter(self.sampler)。如果注释掉while True，DataLoaderIter.__iter__()也迭代完成了，进行单步调试后，next(self.loader)却没有得到值，所以会抛出 StopIteration 异常；当没有注释掉while True时，进行单步调试后，又会跳到Sampler.__iter__()，此时的index=0。

# 打印三遍0-99
class Sampler:def __init__(self, data):self.data = datadef __iter__(self):for index in range(len(self.data)):yield self.data[index]class DataLoaderIter:def __init__(self, loader):# self.sampler = iter(loader.sampler)self.sampler = loader.samplerdef __iter__(self):while True:  # 注释遍历完第一个epoch后就会报错yield from iter(self.sampler)# print("报错")class DataLoader:def __init__(self, sampler):self.sampler = samplerself.loader = iter(DataLoaderIter(self))def __iter__(self):for _ in range(100):yield next(self.loader)# 使用示例
if __name__ == "__main__":sampler = Sampler(range(100))  # 假设的采样器实例data_loader = DataLoader(sampler)epochs = 3for i in range(epcohs):for item in data_loader:print(item)

2 YOLOv9

create_dataloader代码

def create_dataloader(path,imgsz,batch_size,stride,single_cls=False,hyp=None,augment=False,cache=False,pad=0.0,rect=False,rank=-1,workers=8,image_weights=False,close_mosaic=False,quad=False,min_items=0,prefix='',shuffle=False):if rect and shuffle:LOGGER.warning('WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False')shuffle = Falsewith torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDPdataset = LoadImagesAndLabels(path,imgsz,batch_size,augment=augment,  # augmentationhyp=hyp,  # hyperparametersrect=rect,  # rectangular batchescache_images=cache,single_cls=single_cls,stride=int(stride),pad=pad,image_weights=image_weights,min_items=min_items,prefix=prefix)batch_size = min(batch_size, len(dataset))nd = torch.cuda.device_count()  # number of CUDA devicesnw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])  # number of workerssampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)#loader = DataLoader if image_weights else InfiniteDataLoader  # only DataLoader allows for attribute updatesloader = DataLoader if image_weights or close_mosaic else InfiniteDataLoadergenerator = torch.Generator()generator.manual_seed(6148914691236517205 + RANK)return loader(dataset,batch_size=batch_size,shuffle=shuffle and sampler is None,num_workers=nw,sampler=sampler,pin_memory=PIN_MEMORY,collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,worker_init_fn=seed_worker,generator=generator), dataset

在YOLOv9中，当close_mosaic = 15(启用close_moasic)时，读取数据采用的是torch.utils.data.DataLoader，当close_mosaic=0，读取数据采用的是InfiniteDataLoader。

2.1 DataLoader

2.1.1 DataLoader(介绍使用torch.utils.data.DataLoader读取数据)

再次强调，一定要进行调试理解
下面介绍按照章节1.2.3进行理解

读取数据集基本流程

train_loader.__iter__() => DataLoader.__iter__() => _SingleProcessDataLoaderIter(self) => _BaseDataLoaderIter.__next__() => _SingleProcessDataLoaderIter._next_data() => index = self._next_index() => data = self._dataset_fetcher.fetch(index)

获得数据集索引的流程
index = self._next_index() => _BaseDataLoaderIter._next_index() => DataLoader._index_sampler() => BatchSampler.__iter__() => RandomSampler.__iter__()

2.1.1.1 for循环：生成迭代器

对于一个epoch，使用的时for i, (imgs, targets, paths, _) in pbar:，忽略掉tqdm，tqdm只是为了显示读取进度条的作用，即可以当成for i, (imgs, targets, paths, _) in enumerate(train_loader): 。

train_loader = create_dataloader()， create_dataloader使用DataLoader 。

下面代码注释部分是train，使用官方提供的训练命令对应的初始化值。

# loader = DataLoader 
loader(dataset,batch_size=batch_size,  # 32shuffle=shuffle and sampler is None,  # Truenum_workers=nw,  # 8sampler=sampler,  # Nonepin_memory=PIN_MEMORY,  # Truecollate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,  # LoadImagesAndLabels.collate_fnworker_init_fn=seed_worker,  generator=generator), dataset  

按照章节1.2中介绍的，DataLoader 会调用DataLoader .__iter__()返回一个迭代器，从类DataLoader 中的__iter__()代码可以看出，由于self.persistent_workers 默认为False，返回的迭代器是self._get_iterator()。

# class DataLoaderdef __iter__(self) -> '_BaseDataLoaderIter':# When using a single worker the returned iterator should be# created everytime to avoid resetting its state# However, in the case of a multiple workers iterator# the iterator is only created once in the lifetime of the# DataLoader object so that workers can be reusedif 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()

接下来，从下面代码可以看出，当num_workers = 8时，返回的迭代器是_MultiProcessingDataLoaderIter(self)。当num_workers = 0时，返回的迭代器是_SingleProcessDataLoaderIter(self)。

# class DataLoaderdef _get_iterator(self) -> '_BaseDataLoaderIter':if self.num_workers == 0:return _SingleProcessDataLoaderIter(self)else:self.check_worker_number_rationality()return _MultiProcessingDataLoaderIter(self)

DataLoader相对于InfiniteDataLoader最大的区别

使用torch.utils.data.DataLoader读取数据集，每一轮epoch返回的迭代器_MultiProcessingDataLoaderIter(self)或者_SingleProcessDataLoaderIter(self)都对这两个类重新进行了初始化。
使用InfiniteDataLoader，在InfiniteDataLoader.__init__就初始化_MultiProcessingDataLoaderIter(self)或者_SingleProcessDataLoaderIter(self)，在epochs每一轮期间就不需要重新初始化了。

2.1.1.2 for循环：生成迭代器后调用迭代器的__next__()

_MultiProcessingDataLoaderIter是继承于_BaseDataLoaderIter，为了方便理解，令num_workers = 0，也就是返回的迭代器是_SingleProcessDataLoaderIter(self)，这个也是继承_BaseDataLoaderIter。(两个迭代器的实现过程是一样的，只不过，_MultiProcessingDataLoaderIter是多进程读取，代码相对复杂一点)。

然后就是调用_SingleProcessDataLoaderIter的__next__()，在_SingleProcessDataLoaderIter中并没有__next__()，因此调用的是_BaseDataLoaderIter的__next__()。下面的这行代码是返回用来返回数据的

# class _BaseDataLoaderIter: def __next__(self) data = self._next_data()  

self._next_data()是在_SingleProcessDataLoaderIter中重写

class _SingleProcessDataLoaderIter(_BaseDataLoaderIter):def __init__(self, loader):super().__init__(loader)assert self._timeout == 0assert self._num_workers == 0# Adds forward compatibilities so classic DataLoader can work with DataPipes:#   Taking care of distributed shardingif isinstance(self._dataset, (IterDataPipe, MapDataPipe)):# For BC, use default SHARDING_PRIORITIEStorch.utils.data.graph_settings.apply_sharding(self._dataset, self._world_size, self._rank)self._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 StopIterationdata = self._dataset_fetcher.fetch(index)  # may raise StopIterationif self._pin_memory:data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)return data

self._next_index()是返回的数据集的索引，self._dataset_fetcher将返回的索引变为训练需要的数据

# class _BaseDataLoaderIterdef _next_index(self):return next(self._sampler_iter)  # may raise StopIteration

self._sampler_iter是一个迭代器，是在_SingleProcessDataLoaderIter初始化中定义的

# class _BaseDataLoaderIterself._sampler_iter = iter(self._index_sampler)

关键就来到这个self._index_sampler

# class _BaseDataLoaderIterself._index_sampler = loader._index_sampler 

2.1.1.3 BatchSampler

loader是_SingleProcessDataLoaderIter(self)中的self(即是DataLoader类本身), 在DataLoader类中找到下面代码，知_index_sampler = self.batch_sampler

# class DataLoaderdef _index_sampler(self):# The actual sampler used for generating indices for `_DatasetFetcher`# (see _utils/fetch.py) to read data at each time. This would be# `.batch_sampler` if in auto-collation mode, and `.sampler` otherwise.# We can't change `.sampler` and `.batch_sampler` attributes for BC# reasons.if self._auto_collation:return self.batch_samplerelse:return self.sampler

self.batch_sampler就是是batch_sampler，batch_sampler

# class DataLoaderbatch_sampler = BatchSampler(sampler, batch_size, drop_last)
self.batch_sampler = batch_sampler

章节1.2.3中Sampler有一个__iter__()，是通过这个迭代数据的，因此在BatchSampler中也是通过__iter__()迭代数据的。for idx in self.sampler:的含义是获得datasets的idx序号，将其打包成一个batch。

# class BatchSamplerdef __iter__(self) -> Iterator[List[int]]:# Implemented based on the benchmarking in https://github.com/pytorch/pytorch/pull/76951if self.drop_last:sampler_iter = iter(self.sampler)while True:try:batch = [next(sampler_iter) for _ in range(self.batch_size)]yield batchexcept StopIteration:breakelse:batch = [0] * self.batch_sizeidx_in_batch = 0for idx in self.sampler:batch[idx_in_batch] = idxidx_in_batch += 1if idx_in_batch == self.batch_size:yield batchidx_in_batch = 0batch = [0] * self.batch_sizeif idx_in_batch > 0:yield batch[:idx_in_batch]

2.1.1.4 sample

在DataLoader中最关键的是sample和batch_sampler，当shuffle=True，sample=RandomSampler，当shuffle=False时，sample=SequentialSampler。

if shuffle:sampler = RandomSampler(dataset, generator=generator)  # type: ignore[arg-type]
else:sampler = SequentialSampler(dataset)  # type: ignore[arg-type]

因此，对于sample来说，shuffle(bool)决定了采用RandomSampler还是SequentialSampler。

RandomSampler是打乱数据集的idx ，也就是打乱DataLoader读取datasets的顺序。SequentialSampler是不打乱DataLoader读取datasets的顺序。

注意，如果shuffle=True, 也就是采用RandomSampler，不管是DataLoader还是InfiniteDataLoader，每一epoch中的batch中图片idx都是不一样的（前一epoch和后一epoch的同一batch读取的图片是不一样）。

2.2 InfiniteDataLoader

读取数据集基本流程

train_loader.__iter__() => InfiniteDataLoader.__iter__() => next(self.iterator) => _SingleProcessDataLoaderIter(self) => _BaseDataLoaderIter.__next__() => _SingleProcessDataLoaderIter._next_data() => index = self._next_index() => data = self._dataset_fetcher.fetch(index)

获得数据集索引的流程
index = self._next_index() => _BaseDataLoaderIter._next_index() => DataLoader._index_sampler() => _RepeatSampler.__iter__() => iter(self.sampler) => BatchSampler.__iter__() => RandomSampler.__iter__()

细节部分可以按照章节1.2.4理解
获得数据集索引的流程中存在RandomSampler.__iter__()的原因是，初始化_SingleProcessDataLoaderIter(self)时，章节2.1.1.2中loader._index_sampler此时为_RepeatSampler(self.batch_sampler)

class InfiniteDataLoader(dataloader.DataLoader):""" Dataloader that reuses workersUses same syntax as vanilla DataLoader"""def __init__(self, *args, **kwargs):super().__init__(*args, **kwargs)object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))self.iterator = super().__iter__()def __len__(self):return len(self.batch_sampler.sampler)def __iter__(self):for _ in range(len(self)):yield next(self.iterator)class _RepeatSampler:""" Sampler that repeats foreverArgs:sampler (Sampler)"""def __init__(self, sampler):self.sampler = samplerdef __iter__(self):while True:yield from iter(self.sampler)

2.3 YOLOv9 close_mosaic

章节2提到，在YOLOv9中，当close_mosaic = 15(启用close_moasic)时，读取数据采用的是torch.utils.data.DataLoader，当close_mosaic=0，读取数据采用的是InfiniteDataLoader。

2.3.1 YOLOv9 实现close_mosaic流程

在train.py中，在for epoch in range(start_epoch, epochs):中有判断

if epoch == (epochs - opt.close_mosaic):LOGGER.info("Closing dataloader mosaic")dataset.mosaic = False

就是这么的简单，当然，读取数据集只需要采用torch.utils.data.DataLoader即可。

2.3.2 为什么close_mosaic=15的时候，读取数据只能采取DataLoader?

小tip：YOLOv9使用close_mosaic，调用DataLoader读取数据，每一个epoch前都会等待一会，这是正常现象，因为需要重新初始化迭代器。

DataLoader可以

这个其实来源于DataLoader相对于InfiniteDataLoader最大的区别那一小章节，由于DataLoder在每一epoch都需要重新初始化迭代器_SingleProcessDataLoaderIter(self)，类_SingleProcessDataLoaderIter继承了类_BaseDataLoaderIter。在_BaseDataLoaderIte.__init__()中有

# class _BaseDataLoaderIter: def __init__()
self._dataset = loader.dataset# loader.dataset
# class DataLoader: def __init()
self.dataset = dataset

获取数据在章节2.1.1.2中的

# class _SingleProcessDataLoaderIter : def _next_data()
data = self._dataset_fetcher.fetch(index)  # may raise StopIteration

这个获取数据主要是将index传入构建数据集dataset类LoadImagesAndLabels中的__getitem__中

# class LoadImagesAndLabels: def __getitem__()
mosaic = self.mosaic and random.random() < hyp['mosaic']

2.3.1中是将dataset.mosaic = False，这也会让train_loader.dataset.mosaic=False，这个可以调试一下下面代码进行理解，也可以直接在YOLOv9中进行调试，将断点设置在if epoch == (epochs - opt.close_mosaic):这一行，当运行完dataset.mosaic = False之后，train_loader.dataset.mosaic=False。

class Dataset:def __init__(self):self.mosaic = Truedef print(self):print(self.mosaic)class DataLoder:def __init__(self, dataset):self.dataset = datasetdef print(self):print(self.dataset.mosaic)if __name__ == "__main__":dataset = Dataset()train_loader = DataLoder(dataset)train_loader.print()dataset.mosaic = Falsetrain_loader.print()

在将train_loader.dataset.mosaic置为False后，迭代器_SingleProcessDataLoaderIter中的_SingleProcessDataLoaderIter._dataset.mosaic也已经置为False，但是在迭代器_SingleProcessDataLoaderIter(self)初始化的时候，从下面的代码知，已经调用了_SingleProcessDataLoaderIter(self)._dataset。因此需要在每一epoch将迭代器_SingleProcessDataLoaderIter重新进行初始化，将train_loader.dataset.mosaic=False的dataset传入到_SingleProcessDataLoaderIter对_SingleProcessDataLoaderIter._dataset_fetcher进行初始化。在使用close_mosaic之后，重新初始化迭代器后_SingleProcessDataLoaderIter._dataset.mosaic = False，使用self._dataset_fetcher.fetch(index)调用dataset.__getitem__获取数据时的mosaic就为False了。

# class _SingleProcessDataLoaderIter: def __init__()self._dataset_fetcher = _DatasetKind.create_fetcher(
self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last)

InfiniteDataLoader不可以

而对于InfiniteDataLoader而言，其是在create_dataloader()时，就已经在InfiniteDataLoader.__init__()中初始化了self.iterator = super().__iter__() => self.iterator = _SingleProcessDataLoaderIter(self)。

当dataset.mosaic = False，train_loader.iterator._dataset.mosaic确实也置为False(可以按照上一小节DataLoader可以中的断点进行调试)，但是在迭代器_SingleProcessDataLoaderIter(self)初始化的时候，从下面代码可以看出，已经调用了_SingleProcessDataLoaderIter(self)._dataset。因此，self._dataset_fetcher的值是初始化时候的_dataset.mosaic决定的，在epoch中改变的dataset.mosaic可以改变_dataset.mosaic的值，但是改变不了_dataset_fetcher。

# class _SingleProcessDataLoaderIter: def __init__()self._dataset_fetcher = _DatasetKind.create_fetcher(
self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last)

2.4 YOLOv8(version = “8.2.31”) 实现close_mosaic流程

YOLOv8中的InfiniteDataLoader有一点一点不同于YOLOv9

class InfiniteDataLoader(dataloader.DataLoader):"""Dataloader that reuses workers.Uses same syntax as vanilla DataLoader."""def __init__(self, *args, **kwargs):"""Dataloader that infinitely recycles workers, inherits from DataLoader."""super().__init__(*args, **kwargs)object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))self.iterator = super().__iter__()def __len__(self):"""Returns the length of the batch sampler's sampler."""return len(self.batch_sampler.sampler)def __iter__(self):"""Creates a sampler that repeats indefinitely."""for _ in range(len(self)):yield next(self.iterator)def reset(self):"""Reset iterator.This is useful when we want to modify settings of dataset while training."""self.iterator = self._get_iterator()class _RepeatSampler:"""Sampler that repeats forever.Args:sampler (Dataset.sampler): The sampler to repeat."""def __init__(self, sampler):"""Initializes an object that repeats a given sampler indefinitely."""self.sampler = samplerdef __iter__(self):"""Iterates over the 'sampler' and yields its contents."""while True:yield from iter(self.sampler)

YOLOv8不管是使用close_mosaic还是不适用close_mosaic，使用的数据加载器都是InfiniteDataLoader，从下面代码可知，在将self.train_loader.dataset.mosaic = False，后又重新调用了self.train_loader.reset() => self._get_iterator()重新初始化了一个迭代器_SingleProcessDataLoaderIter(self)，此时的_SingleProcessDataLoaderIter._dataset_fetcher的参数_SingleProcessDataLoaderIter._dataset.mosaic就是False。

# class BaseTrainer: def _do_train()
if epoch == (self.epochs - self.args.close_mosaic):self._close_dataloader_mosaic()self.train_loader.reset()def _close_dataloader_mosaic(self):"""Update dataloaders to stop using mosaic augmentation."""if hasattr(self.train_loader.dataset, "mosaic"):self.train_loader.dataset.mosaic = Falseif hasattr(self.train_loader.dataset, "close_mosaic"):LOGGER.info("Closing dataloader mosaic")self.train_loader.dataset.close_mosaic(hyp=self.args)

3 参考链接

1. VScode如何Debug(调试)进入标准库文件/第三方包源码