数据加载

yolov5的数据加载部分由create_dataloader函数实现（位于utils/datasets.py），其中关于数据增强和加载的部分主要由LoadImagesAndLabels和InfiniteDataLoader负责，并基于torch_distributed_zero_first(rank)进行不同进程之间的数据同步。

数据同步

在yolov5的模型训练中涉及了多进程并行运算。其中，主进程实现数据的预读取并缓存，然后其它子进程则从缓存中读取数据并进行一系列运算。为了完成数据的正常同步，yolov5中基于torch.distributed.barrier()函数实现了上下文管理器torch_distributed_zero_first：

@contextmanager
def torch_distributed_zero_first(local_rank: int):
    """
    Decorator to make all processes in distributed training wait for each local_master to do something.
    """
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()
    yield
    if local_rank == 0:
        torch.distributed.barrier()

torch_distributed_zero_first使用方式：

with torch_distributed_zero_first(rank):
     dataset = LoadImagesAndLabels(......)

rank表示当前的进程号，主进程由编号0表示，子进程则由编号1、2、3...等表示。上述代码的运行逻辑：

1. 进入torch_distributed_zero_first(rank)上下文作用域；
2. 判断当前进程号local_rank是否为-1或0，如果不是则说明为子进程，运行torch.distributed.barrier()等待主进程的数据处理完毕，如果是则当前为主进程，不需要等待；
3. yield后，运行with作用域范围内的代码；
4. 作用域范围内代码运行完成后，继续yield的后续操作，判断当前进程号是否为0（即是否为主进程），如果是，则运行torch.distributed.barrier()可以解开其它子进程的阻塞。

注意 对于torch.distributed.barrier()函数的作用可以参考https://stackoverflow.com/questions/59760328/how-does-torch-distributed-barrier-work进行理解。

数据增强

数据增强相关的方法在LoadImagesAndLabels类中实现。

Rectangular Training

通常YOLO系列网络的输入都是预处理后的方形图像数据，如416 * 416、608 * 608。当原始图像为矩形时，会将其填充为方形（如下图：方形输入），但是填充的灰色区域其实就是冗余信息，不论是在训练还是推理阶段，这些冗余信息都会增加耗时。

方形输入

为了减少图像的冗余数据，输入图像由方形改为矩形（如下图：矩形输入）：将长边resize为固定尺寸（如416），短边按同样比例resize，然后把短边的尺寸尽量少地填充为32的倍数。

矩形输入

这种方法在推理阶段称为矩形推理（Rectangular Inference），在训练阶段则称为矩形训练（Rectangular Training）。推理阶段直接对图像进行resize和pad就行，但是训练阶段输入的是一个批次的图像集合，需要保持批次内的图像尺寸一致，因此处理逻辑相对复杂一些。

代码：

        if self.rect:
            # Sort by aspect ratio
            # 首先根据高宽比排序，就可以保证每个batch内的图像高宽比相近。
            s = self.shapes  # wh
            ar = s[:, 1] / s[:, 0]  # aspect ratio    高/宽
            irect = ar.argsort()
            self.img_files = [self.img_files[i] for i in irect]
            self.label_files = [self.label_files[i] for i in irect]
            self.labels = [self.labels[i] for i in irect]
            self.shapes = s[irect]  # wh
            ar = ar[irect]

            # Set training image shapes
            shapes = [[1, 1]] * nb    # hw
            for i in range(nb):
                ari = ar[bi == i]
                mini, maxi = ari.min(), ari.max()
                if maxi < 1:    # 高宽比最大值都小于1，则说明batch内的图全都是高小于宽
                    shapes[i] = [maxi, 1]   # 设置宽为固定比例1，高的比例为maxi
                elif mini > 1:   # 高宽比最小值都大于1，说明batch内的图都是高>宽
                    shapes[i] = [1, 1 / mini]    # 设置高为固定比例1，宽的比例为1 / mini

运行逻辑：

1. 根据数据集中所有图像的shape计算高宽比ar；
2. 对长宽比ar进行argsort，即对ar内的元素进行排序（升序），并针对排序后的元素对应取得其在ar中的索引，构成索引序列irect；
3. 根据索引序列irect取得排序后的self.img_files、self.label_files、self.labels 、self.shapes和ar；
4. 初始化nb（即batch数量）个shape为[1,1]，组成shapes；
5. 从ar中对应取出每个batch的高宽比列表ari，取其中的最大、最小值；
6. 如果当前batch的高宽比最大值小于1，则将shapes内该batch对应的值设为[maxi, 1]，而如果最大值<=1且最小值>1，则设置为[1, 1 / mini]，如果都不符合默认[1,1]。

参考：https://github.com/ultralytics/yolov3/issues/232

Mosaic

未完待续