YOLOv10s训练代码解析4：构建Dataset和Dataloader

本专栏会手把手带你从源码了解YOLOv10(后续会陆续介绍YOLOv8、RTDETR等模型)，尽可能地完整介绍整个算法，这个专栏会持续创作与更新，大家如果想要本文PDF和思维导图，后台私信我即可(创作不易，不喜勿喷)，大家如果发现任何错误和需要修改的地方都可以私信我，我会统一修改。

模型和训练器加载和初始化后，开始训练(ultralytics/engine/model.py)，之后选择单卡或多卡训练，继续调用ultralytics/engine/trainer.py的第323行_do_train函数

调用ultralytics/engine/trainer.py文件的第236行_setup_train函数：检查AMP、图像尺寸、加载训练集的dataloaders等

(重点) 继续调用ultralytics/models/yolo/detect/train.py的第45行get_dataloader函数，加载训练数据集dataset和构建dataloaders；调用ultralytics/models/yolo/detect/train.py的第33行build_dataset函数；再继续调用ultralytics/data/build.py的第85行build_yolo_dataset函数；后续调用ultralytics/data/dataset.py的第35行YOLODataset类的__init__函数，YOLODataset继承自BaseDataset(ultralytics/data/base.py)，再继续调用初始化(第73行self.im_files返回的是一个列表，每个元素是图片的路径；第74行self.labels返回的是一个列表，每个元素的标签的路径)

调用ultralytics/data/dataset.py中的build_transformers方法构建数据集的数据增强策略；具体的数据增强策略包括：ultralytics/data/augument.py文件中的Mosaic(p=1)、CopyPaste(p=0)、RandamPerspective、MixUp(p=0)、Albumentation(Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01), CLAHE(p=0.01, clip_limit=(1, 4.0), tile_grid_size=(8, 8)))、RandomHSV、RandomFlip(direction="vertical")、RandomFlip(direction="horizontal")，具体如下所示

在第五幅图初始化构建好dataset之后，继续调用第55行来构建dataloader；ultralytics/data/build.py的第29行的InfiniteDataLoader类的初始化__init__方法，主要用于创建一个可以无限循环的数据加载器，在训练过程中不断地从数据集中采样，而无需在每次迭代结束后重新初始化 DataLoader（具体实现参考InfiniteDataLoader和_RepeatSampler类）

第四幅图加载Dataset和Dataloader之后，继续调用ultralytics/engine/trainer.py:的加载验证集的Dataset和Dataloader（与训练集的dataloader类似，略过）、初始化验证器validator(ultralytics/models/yolov10/train.py的第297行获取validator，继续调用ultralytics/models/yolov10/val.py文件的YOLOv10DetectionValidator类的__init__方法，YOLOv10DetectionValidator类继承自DetectionValidator类，调用DetectionValidator类的__init__方法(ultralytics/models/yolo/detect/val.py)，DetectionValidator又继承自BaseValidator，调用BaseValidator类的__init__方法(ultralytics/engine/validator.py))和指标metrics(metric_keys: ['metrics/precision(B)', 'metrics/recall(B)', 'metrics/mAP50(B)', 'metrics/mAP50-95(B)', 'val/box_om', 'val/cls_om', 'val/dfl_om', 'val/box_oo', 'val/cls_oo', 'val/dfl_oo'])等

在初始化validator和metrics之后，ultralytics/engine/trainer.py继续初始化优化器、学习率等

第三幅图_do_train函数初始化训练配置之后，ultralytics/engine/trainer.py：(自己训练时，nb为8，nw是warmup iterations，由于数据集太小这个值是100，base_idx是关闭mosaic用到的，值为720=8*(100-10))；后续开始训练；ultralytics/engine/trainer.py: warmup调整学习率、动量等参数，之后进入真正的数据处理、前向传播和梯度更新

对于上图的batch从Dataloader取数据，ultralytics/data/base.py的__getitem__函数取数据（self.rect将图像尺寸变换成32的整数倍）

batch信息：im_files(图像路径，长度为batch)、ori_shape(图像原始尺寸，长度为batch)、resized_shape(图像resize后的尺寸，长度为batch)、img(tensor维度为[batch, 3, 640, 640])、cls(tensor维度为[Nobj])、bboxes(tensor维度为[Nobj, 4])、batch_idx(每个目标对象对应于第几张图片, 0~batch-1)

ultralytics/data/base.py文件的load_image函数读取图像，(h0, w0)为原始图像尺寸，im.shape[:2]为变换后的图像尺寸(最长边到640)

ultralytics/data/dataset.py文件的update_labels_info函数更新目标边框信息（cache的self.labels里面存储了目标边框信息）

ultralytics/data/augment.py的数据增强(下一章节重点介绍)

这部分关于构建Dataset和Dataloader的整体流程如下图所示，需要原图和思维导图的朋友关注我私信获取。