代码解读 1：一起读 Ultralytics 库工程化的 YOLO 代码 (rt-detr, v3, v5, v6, v8, v9, v10, v11)

君名余曰正则

已于 2024-12-17 10:59:12 修改

阅读量1.7k

点赞数 21

分类专栏：计算机视觉文章标签： YOLO Ultralytics 计算机视觉深度学习

于 2024-12-16 16:03:22 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_41782172/article/details/144196868

版权

计算机视觉专栏收录该内容

5 篇文章

订阅专栏

代码解读系列文章目录

2024年12月16日

结束了在学校做研究的生活进入了业界，更多将接触工程化的内容。YOLO 作为目标检测领域出众的模型，其中 ultralytics 库的工程化做的也相当出色，看大佬们的代码结构，各个模块的设计也能够学到很多。官方的 ultralytics 库地址 https://github.com/ultralytics/ultralytics

考虑到很多人可能和我一样，工程化代码接触的并不多，因为研究需要有一点跨领域刚接触 ultralytics 库，如何快速了解 ultralytics 库结构并根据实际需要添加代码。

因此开启这样一个专栏来一起学习、讨论相关工程化的代码，以及模型的创新点设计。此外，也分享、讨论在实际使用过程中遇到的一些优化问题和可能可行的解决方法。

在这个专栏中呢，我将尽可能详细的介绍整个 ultralytics 工程化项目的代码结构。文章目前考虑主要面向有一定计算机视觉基础和神经网络基础的学习者，模型方面打算只做简单的介绍。

前言

对于一个新开的专栏来说，希望读者朋友们提供建议，一起讨论。自己在代码工程化部分接触的还不算多，也在学习阶段，可能有不对的地方还恳请大佬们指正。

第一篇文章先从 ultralytics 项目的整体架构开始，主要介绍文件路径、各个文件夹代表的功能模块以及一些重要的 .py 文件。

一、文件路径及代表的功能模块

├── docker                                  # 官方的 docker 版本（不用 docker 的话不用管）
├── docs                                  # 官方文档
├── examples                                  # 官方实例
├── test                                    # 测试文件
└── ultralytics                             # 核心的模型代码文件和其他工具组件
     ├── assets                             # 感觉好像没啥用的
     ├── cfg                                # --- 参数 ---
     │    ├── datasets                      # 数据集的参数
     │    ├── models                        # 各类模型的网络结构参数 rt-detr,v3,v5,v6,v8,v9,v10,v11
     │    ├── trackers                      # ？ 部分参数文件
     │    └── default.yaml                  # ！ 训练器默认会加载的参数
     ├── data                               # --- 数据处理 ---
     │    ├── explorer                      # 可视化
     │    ├── scripts                       # 获取数据的脚本文件
     │    ├── augment.py                    # ！ 所有的数据增强类和实现
     │    └── ***.py                        # 各类数据读取器、数据处理等相关代码
     ├── engine                             # --- 最底层的类定义 ---
     │    ├── model.py                      # 模型基类
     │    ├── trainer.py                    # 训练器基类
     │    └── ***.py                        # 其他各种基类
     ├── hub                                # ？
     ├── models                             # --- 外层的类定义 ---
     │    ├── fastsam                       # fast SAM 模型
     │    ├── nas                           # NAS 模型
     │    ├── rtdetr                        # RT-DETR 模型
     │    ├── sam                           # SAM 模型
     │    │    └── modules                  # SAM中用到的网络结构的实现，encoder、decoder等
     │    ├── utils                         # 各网络中公用的基础组件
     │    └── yolo                          # YOLO 模型
     │         ├── classify                 # 分类任务
     │         ├── detect                   # 检测任务
     │         ├── obb                      # Oriented Bounding Box 3D包围盒
     │         ├── pose                     # 姿态估计
     │         └── segment                  # 分割任务
     ├── nn                                 # --- 公用的神经网络组件 ---
     │    └── modules                       # 公用的基础网络结构的实现
     ├── solutions                          # --- 一些偏专用的解决方法 ---
     ├── trackers                           # ？
     │    └── utils                         # 该任务中公用的基础组件
     └── utils                              # --- 在架构各阶段用到的公用的基础组件 ---
          └── callbacks                     # ？ 回调函数、可视化等基础组件

二、代码运行逻辑

mytrain.py 中的 YOLO(model_yaml, task='detect')，模型初始化
- 先基类 engine/model.py 中的 class Model(nn.Module) 的初始化 def __init__
  - 被加载的内容
    - 默认的参数文件 cfg/default.yaml
  - 由 task 任务类型确定的子类模型 models/yolo/model.py 赋值给 self.model,整个模型初始化基本完成
    - 这个过程会调用 nn/tasks.py 中的 class DetectionModel 的初始化 def __init__
      - 被加载的内容
        模型配置参数文件 yolov8s.yaml
- 再执行子类 models/yolo/model.py 中的 def __init__
mytrain.py 中的 model.train
- 进入第 3 步
engine/model.py 中的 def train,模型基类的公用 train 函数
- 训练器实例初始化。由 task 任务类型确定的训练器类型为 models/yolo/detect/train.py 中的 class DetectionTrainer(BaseTrainer)
  - 只有基类 engine/trainer.py 中的 class BaseTrainer 的初始化 def __init__
  - 被加载的内容
    - 作为参数的自定义的数据配置文件 soldier.yaml
- 训练器初始化。调用 models/yolo/detect/train.py 中的 def get_model
  - 通过 self.model 中的配置参数来加载模型
- 训练器启动训练 self.trainer.train()。调用训练器基类 engine/trainer.py 的 def train
  - 进入第 4 步
engine/trainer.py 中的 def train
- 其中主要做 GPU 判断和分布式判断和设置
  - 如果是分布式的话，通过生成的训练的 .py 临时文件和命令启动训练
- 单个进程的训练走的是 engine/trainer.py 中的 def _do_train
  - 进入第 5 步
engine/trainer.py 中的 def _do_train
- 先初始化训练器的各组件。通过 engine/trainer.py 中的 def _setup_train
  - 进入第 6 步
- 计算 warmup 热身阶段迭代次数
- 训练只剩下最后多少 epoch 时停用 mosaic 图像增强 close_mosaic 参数控制
- 进入训练 epoch 循环
  - 从数据读取器中获取本轮的 batch 数据
  - 每次 batch 个数据训练的循环
    - warmup 热身阶段判断
      - 在初始训练阶段逐步调整学习率和动量，避免大梯度带来的训练振荡或数值不稳定
    - 数据前向传播
      self.loss, self.loss_items = self.model(batch)
      - 通过 nn/tasks.py 中的 class DetectionModel(BaseModel)
        进入第 12 步
    - 梯度反向传播
    - 执行优化器
      - 每当距离上次更新参数后，累积的梯度步数达到 self.accumulate 时，更新参数
    - 记录日志
  - 计算各项指标并生成日志，保存日志
  - 更新学习率
- 训练结束
  - 打印总共完成的 epoch 数和训练耗时
  - 使用最优模型（如 best.pt）进行最后一次评估
  - 生成训练指标的图表（plots 默认为 True）
  - 释放 GPU 显存
engine/trainer.py 中的 def _setup_train
- 冻结层
- 是否使用混合精度来加速训练，并减少显存占用
- 根据模型的最大步幅确定图像的最小尺寸
- 训练数据加载器初始化。调用 models/yolo/detect/train.py 中的 def get_dataloader
  - 进入第 7 步
- 优化器设置
- 学习率调度器和早停设置
models/yolo/detect/train.py 中的 def get_dataloader
- 构建数据类的对象
  - data/build.py 中的 def build_yolo_dataset
    - 训练模式下默认进行数据增强 True
  - data/dataset.py 中的 class YOLODataset(BaseDataset) 初始化 def __init__
    - 该初始化只是做了几步参数判断，然后是其基类 data/base.py 中的 class BaseDataset(Dataset) 初始化 def __init__
    - 进入第 8 步
- 构建数据读取器类的对象
  - data/build.py 中的 def build_dataloader
    - 进入第 11 步
data/base.py 中的 class BaseDataset(Dataset) 初始化 def __init__
- 参数赋值，其中有部分重要的参数
  - cache：是否缓存图像（默认为 False）
  - augment：是否进行图像增强（训练模式为 True，其他默认为 False）
  - rect：是否采用矩阵填充模式（默认不启用）
- 调用类函数 def build_transforms
  - 具体实现在子类 class YOLODataset(BaseDataset) 中的 def build_transforms
    - 进入第 9 步
class YOLODataset(BaseDataset) 中的 def build_transforms
- 如果启用了矩形填充模式（rect=True），则关闭某些增强功能（如 Mosaic 和 Mixup）
- 数据增强函数
  - 调用 data/aygment.py 中的 def v8_transforms
    - 进入第 10 步

data/aygment.py 中的 def v8_transforms

直接放对应的代码，几乎每种增强方法的实现都有对应的函数

def v8_transforms(dataset, imgsz, hyp, stretch=False):
	"""Convert images to a size suitable for YOLOv8 training."""
    # 预处理
    pre_transform = Compose(
        [
            Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic),     # 用于将多张图像随机拼接为一张（Mosaic 数据增强，mosaic 概率默认为 1）
            CopyPaste(p=hyp.copy_paste),                    # 对图像进行对象级增强，即将目标从一张图像复制并粘贴到另一张图像上（copy_paste 概率默认为 0）
            RandomPerspective(                              # 随机执行透视变换，包括旋转、平移、缩放、剪切等操作
                degrees=hyp.degrees,                        # 旋转角度范围（默认 0）
                translate=hyp.translate,                    # 平移范围（默认 0.1）
                scale=hyp.scale,                            # 缩放范围（默认 0.5）
                shear=hyp.shear,                            # 剪切范围（默认为 0）
                perspective=hyp.perspective,                # 透视效果范围（默认为 0）
                pre_transform=None if stretch else LetterBox(new_shape=(imgsz, imgsz)), # 前置处理操作，如果 stretch=False，则使用 LetterBox 进行缩放和填充（默认为 False）
            ),
        ]
    )

    # 处理关键点增强（只有 task == "pose" 启用）
    flip_idx = dataset.data.get("flip_idx", [])  # for keypoints augmentation，用于指定关键点的左右对称关系（如人脸的左右眼）
    if dataset.use_keypoints:
        kpt_shape = dataset.data.get("kpt_shape", None)
        # 如果 flip_idx 为空但启用了左右翻转（fliplr > 0.0），会发出警告并禁用翻转
        if len(flip_idx) == 0 and hyp.fliplr > 0.0:
            hyp.fliplr = 0.0
            LOGGER.warning("WARNING ⚠️ No 'flip_idx' array defined in data.yaml, setting augmentation 'fliplr=0.0'")
        # 如果 flip_idx 的长度与关键点的数量不匹配，则抛出异常
        elif flip_idx and (len(flip_idx) != kpt_shape[0]):
            raise ValueError(f"data.yaml flip_idx={flip_idx} length must be equal to kpt_shape[0]={kpt_shape[0]}")

    # 构建最终转换迭代器
    return Compose(
        [
            pre_transform,
            MixUp(dataset, pre_transform=pre_transform, p=hyp.mixup),               # 混合增强，将两张图像以随机比例叠加，生成新的训练样本（mixup 默认为 0）
            Albumentations(p=1.0),                                                  # 使用 Albumentations 库的增强功能（一个高级图像增强库），对图像执行更多随机操作
            RandomHSV(hgain=hyp.hsv_h, sgain=hyp.hsv_s, vgain=hyp.hsv_v),           # 随机调整图像的色相（0.015）、饱和度（0.7）和亮度（0.4）
            RandomFlip(direction="vertical", p=hyp.flipud),                         # 垂直翻转
            RandomFlip(direction="horizontal", p=hyp.fliplr, flip_idx=flip_idx),    # 水平翻转，结合 flip_idx 处理关键点对称关系
        ]
    )  # transforms

data/aygment.py 中实现图像增强的类，这些类中的 def __call__ 函数在数据生成器中会被自动调用
- class BaseMixTransform 基类
- class Mosaic(BaseMixTransform)
- class CopyPaste
- class RandomPerspective
- class MixUp(BaseMixTransform)
- class Albumentations
- class RandomHSV
- class RandomFlip

data/build.py 中的 def build_dataloader
nn/tasks.py 中的 class DetectionModel(BaseModel)
- 暂时不确定各个函数执行的逻辑
- 大概率会通过专门的 Loss 计算类来处理，通过 utils/loss.py 中的 class v8DetectionLoss 默认调用函数 def __call__
  - 进入第 13 步
utils/loss.py 中的 class v8DetectionLoss 默认调用函数 def __call__

三、如何使用 Ultralytics 训练 YOLO 模型

官方的介绍中，使用 ultralytics 训练 YOLO 十分的简单。其中要注意的参数有两个

模型文件参数，model = YOLO("yolo11n.pt") 这里采用了官方默认的模型名称，就算你没有这个模型文件，官方也会为你下载在imageNet上训练的预训练模型
数据集的说明文件 data="coco8.yaml"，对于标准开源数据集的 yaml 文件官方都已经为我们准备好了，这类 yaml 文件可以在 ultralytics/cfg/datasets 文件夹中找到

需要注意的是：官方给的这个例子采用的参数大部分都是默认的，默认的参数文件是 ultralytics/cfg/default.yaml，在下一篇文章中将来详细介绍这些参数

from ultralytics import YOLO

# 加载模型，如果没有这个模型文件，ultralytics 会自动下载预训练的模型
model = YOLO("yolo11n.pt")

# 训练模型
train_results = model.train(
    data="coco8.yaml",  # 数据集 YAML 路径
    epochs=100,  # 训练轮次
    imgsz=640,  # 训练图像尺寸
    device="cpu",  # 运行设备，例如 device=0 或 device=0,1,2,3 或 device=cpu
)

# 评估模型在验证集上的性能
metrics = model.val()

# 在图像上执行对象检测，需要换成自己的路径
results = model("path/to/image.jpg")
results[0].show()

# 将模型导出为 ONNX 格式
path = model.export(format="onnx")  # 返回导出模型的路径

使用自己的数据集训练，只需要将数据集的说明文件替换成你自己的 .yaml 文件，文件至少需要包括以下内容

# 训练集
train: "（数据集的路径）/train/images"
# 验证集，用于给你每一轮的模型评分
val: "（数据集的路径）/valid/images"

# 标签，及对应的名称
names:
  0: S		# 0 是打的标签，S 是这个标签对应的实际名称
  1: P
  
# 类别数，这里 0，1 两个类别
nc: 2

四、小结

本章主要介绍了 Ultralytics 库的架构，初步感受大佬们的代码架构，后续再深入讲解重要参数和代码文件，以及如何自己微调模型等。新的专题喜欢大家多多支持，互相学习。