内容目录

前言

【课程总结】Day11(下):YOLO的入门使用一节中,我们已经了解YOLO的使用方法,使用过程非常简单,训练时只需要三行代码:引入YOLO,构建模型,训练模型;预测时也同样简单,只需要两行代码:引入YOLO,预测图像即可。以上过程简单主要是ultralytics的代码库已经做了封装,使得使用者集中精力在模型训练和预测上。

为了更加深入了解YOLO的实现原理,本章内容将对YOLO的工程结构、模型构建过程、模型训练过程尝试深入探究。

YOLO项目的工程结构

|- ultralytics/
    |- assets/    # 存放项目中使用的资源文件,如图像、样本数据等。
    |- cfg/     # 存放模型配置文件,包括不同模型的配置信息,如网络结构、超参数等。
    |- data/    # 存放数据集文件和数据处理相关的代码。
    |- engine/  # 存放训练和推理引擎的代码,包括训练、测试、评估等功能的实现。
    |- hub/     # 存放模型库相关的代码和模型文件,用于快速调用和使用预训练模型。
    |- models/  # 存放模型的定义和实现代码,包括不同模型的网络结构和相关函数。
    |- nn/      # 存放神经网络模块的代码,包括各种层次的定义和实现。
    |- solutions/ # 存放解决方案相关的代码和实现,用于特定问题或任务的解决方案。
    |- trackers/  # 存放目标跟踪相关的代码和实现,包括目标追踪算法的实现。
    |- utils/   # 存放通用的工具函数和辅助函数,用于项目中的各种功能和任务。

进一步查看cfg目录的内容如下:

- ultralytics
  |- cfg
    |- datasets  # 数据集处理和加载相关文件
    |- default.yaml  # 默认配置信息文件
    |- models  # 包含不同模型结构的配置文件
      |- yolov8-cls-resnet101.yaml  # 定义 YOLOv8 模型结构的配置文件(ResNet-101 版本)
      |- yolov8-cls-resnet50.yaml  # 定义 YOLOv8 模型结构的配置文件(ResNet-50 版本)
      |- yolov8-cls.yaml  # 定义 YOLOv8 模型结构的配置文件
      |- ...  # 其他 YOLOv8 模型版本的配置文件
    |- trackers  # 目标跟踪相关文件

YOLO的yaml文件解析

yaml文件内容

以yolov8.yaml为例,其内容如下:

nc: 80 # 类别数目,nc代表"number of classes",即模型用于检测的对象类别总数。
scales: # 模型复合缩放常数,例如 'model=yolov8n.yaml' 将调用带有 'n' 缩放的 yolov8.yaml
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024]  # YOLOv8n概览:225层, 3157200参数, 3157184梯度, 8.9 GFLOPs
  s: [0.33, 0.50, 1024]  # YOLOv8s概览:225层, 11166560参数, 11166544梯度, 28.8 GFLOPs
  m: [0.67, 0.75, 768]   # YOLOv8m概览:295层, 25902640参数, 25902624梯度, 79.3 GFLOPs
  l: [1.00, 1.00, 512]   # YOLOv8l概览:365层, 43691520参数, 43691504梯度, 165.7 GFLOPs
  x: [1.00, 1.25, 512]   # YOLOv8x概览:365层, 68229648参数, 68229632梯度, 258.5 GFLOPs

# YOLOv8.0n 骨干层
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]]   # 0-P1/2 第0层,-1代表将上层的输入作为本层的输入。第0层的输入是640*640*3的图像。Conv代表卷积层,相应的参数:64代表输出通道数,3代表卷积核大小k,2代表stride步长。
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4 第1层,本层和上一层是一样的操作(128代表输出通道数,3代表卷积核大小k,2代表stride步长)
  - [-1, 3, C2f, [128, True]]   #        第2层,本层是C2f模块,3代表本层重复3次。128代表输出通道数,True表示Bottleneck有shortcut。
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8 第3层,进行卷积操作(256代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为80*80*256(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样),特征图的长宽已经变成输入图像的1/8。
  - [-1, 6, C2f, [256, True]]   #        第4层,本层是C2f模块,可以参考第2层的讲解。6代表本层重复6次。256代表输出通道数,True表示Bottleneck有shortcut。经过这层之后,特征图尺寸依旧是80*80*256。
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16第5层,进行卷积操作(512代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为40*40*512(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样),特征图的长宽已经变成输入图像的1/16。
  - [-1, 6, C2f, [512, True]]   #        第6层,本层是C2f模块,可以参考第2层的讲解。6代表本层重复6次。512代表输出通道数,True表示Bottleneck有shortcut。经过这层之后,特征图尺寸依旧是40*40*512。
  - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32第7层,进行卷积操作(1024代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为20*20*1024(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样),特征图的长宽已经变成输入图像的1/32。
  - [-1, 3, C2f, [1024, True]]  #        第8层,本层是C2f模块,可以参考第2层的讲解。3代表本层重复3次。1024代表输出通道数,True表示Bottleneck有shortcut。经过这层之后,特征图尺寸依旧是20*20*1024。
  - [-1, 1, SPPF, [1024, 5]] # 9         第9层,本层是快速空间金字塔池化层(SPPF)。1024代表输出通道数,5代表池化核大小k。结合模块结构图和代码可以看出,最后concat得到的特征图尺寸是20*20*(512*4),经过一次Conv得到20*20*1024。

# YOLOv8.0n 头部层
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']] # 第10层,本层是上采样层。-1代表将上层的输出作为本层的输入。None代表上采样的size(输出尺寸)不指定。2代表scale_factor=2,表示输出的尺寸是输入尺寸的2倍。nearest代表使用的上采样算法为最近邻插值算法。经过这层之后,特征图的长和宽变成原来的两倍,通道数不变,所以最终尺寸为40*40*1024。
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4 第11层,本层是concat层,[-1, 6]代表将上层和第6层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是40*40*1024,第6层的输出是40*40*512,最终本层的输出尺寸为40*40*1536。
  - [-1, 3, C2f, [512]]  # 12 第12层,本层是C2f模块,可以参考第2层的讲解。3代表本层重复3次。512代表输出通道数。与Backbone中C2f不同的是,此处的C2f的bottleneck模块的shortcut=False。
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']] # 第13层,本层也是上采样层(参考第10层)。经过这层之后,特征图的长和宽变成原来的两倍,通道数不变,所以最终尺寸为80*80*512。
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3 第14层,本层是concat层,[-1, 4]代表将上层和第4层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是80*80*512,第6层的输出是80*80*256,最终本层的输出尺寸为80*80*768。
  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)       第15层,本层是C2f模块,可以参考第2层的讲解。3代表本层重复3次。256代表输出通道数。经过这层之后,特征图尺寸变为80*80*256,特征图的长宽已经变成输入图像的1/8。
  - [-1, 1, Conv, [256, 3, 2]] #                 第16层,进行卷积操作(256代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为40*40*256(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样)。
  - [[-1, 12], 1, Concat, [1]]  # cat head P4    第17层,本层是concat层,[-1, 12]代表将上层和第12层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是40*40*256,第12层的输出是40*40*512,最终本层的输出尺寸为40*40*768。
  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)     第18层,本层是C2f模块,可以参考第2层的讲解。3代表本层重复3次。512代表输出通道数。经过这层之后,特征图尺寸变为40*40*512,特征图的长宽已经变成输入图像的1/16。
  - [-1, 1, Conv, [512, 3, 2]] #                 第19层,进行卷积操作(512代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为20*20*512(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样)。
  - [[-1, 9], 1, Concat, [1]]  # cat head P5     第20层,本层是concat层,[-1, 9]代表将上层和第9层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是20*20*512,第9层的输出是20*20*1024,最终本层的输出尺寸为20*20*1536。
  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)     第21层,本层是C2f模块,可以参考第2层的讲解。3代表本层重复3次。1024代表输出通道数。经过这层之后,特征图尺寸变为20*20*1024,特征图的长宽已经变成输入图像的1/32。
  - [[15, 18, 21], 1, Detect, [nc]]  # Detect(P3, P4, P5) 第20层,本层是Detect层,[15, 18, 21]代表将第15、18、21层的输出(分别是80*80*256、40*40*512、20*20*1024)作为本层的输入。nc是数据集的类别数。

yaml文件解析

上述文件包含了 YOLOv8 模型的配置信息,其中包括了模型的类别数目、模型复合缩放常数、主干网络结构和头部网络结构。具体来说:

  • nc 表示模型的类别数目为 1000。
  • scales 包含了不同缩放常数对应的模型结构参数。
    • 子参数:n, s, m, l, x表示不同的模型尺寸,每个尺寸都有对应的depth(深度)、width(宽度)和max_channels(最大通道数)。
    • depth: 表示深度因子,用来控制一些特定模块的数量的,模块数量多网络深度就深;
    • width: 表示宽度因子,用来控制整个网络结构的通道数量,通道数量越多,网络就看上去更胖更宽;
    • max_channels: 最大通道数,为了动态地调整网络的复杂性。在 YOLO 的早期版本中,网络中的每个层都是固定的,这意味着每个层的通道数也是固定的。但在 YOLOv8 中,为了增加网络的灵活性并使其能够更好地适应不同的任务和数据集,引入了 max_channels 参数。
  • backbone 定义了 YOLOv8.0n 模型的主干网络结构,包括了卷积层、C2f 模块等。
    • from(来自)
      • 这个字段表示当前层连接到的上一层的索引。通常,-1 表示连接到上一层,0 表示连接到输入数据。
      • 例如,[-1, 1, Conv, [64, 3, 2]] 表示当前层连接到上一层,即前一层的输出作为当前层的输入。
    • repeats(重复次数)
      • 这个字段表示当前层的模块(module)被重复使用的次数。
      • 例如,[-1, 3, C2f, [128, True]] 表示当前层的模块 C2f 会被重复使用 3 次。
    • module(模块)
      • 这个字段表示当前层使用的模块类型,如 Conv(卷积层)、C2f 等。
      • 例如,[-1, 1, Conv, [64, 3, 2]] 中的 Conv 表示当前层使用的是卷积层。
    • args(参数)
      • 这个字段包含了当前层模块的参数,例如卷积层的通道数、卷积核大小、步长等。
      • 例如,[-1, 1, Conv, [64, 3, 2]] 中的 [64, 3, 2] 表示卷积层的通道数为 64,卷积核大小为 3,步长为 2。

YOLOv8 模型结构

上述yaml定义的模型结构,使用图示显示如下:

主要组成部分:

  • Backbone(主干网络)
    主干网络是模型的基础,负责从输入图像中提取特征。这些特征是后续网络层进行目标检测的基础。
  • Head(头部网络)
    头部网络是目标检测模型的决策部分,负责产生最终的检测结果。
  • ConvModule
    包含卷积层、BN(批量归一化)和激活函数(如SiLU),用于提取特征。
  • DarknetBottleneck:
    通过residual connections(残差结构)增加网络深度,同时保持效率。
  • CSP Layer:
    CSP结构的变体,通过部分连接来提高模型的训练效率。
  • Concat:
    特征图拼接,用于合并不同层的特征。
  • Upsample:
    上采样操作,增加特征图的空间分辨率。

主要过程为:
1.多层卷积:图片输入到主干网络,经过P1-P5层卷积后,在第9层通过SPPF模块进行特征提取。
2.上采样和连接:经过SPPF后进行Upsample上采样操作,然后与第6层进行concat操作,再进行C2f模块特征提取。
类似的:

  • 在14层与第4层进行concat操作,再进行C2f模块特征提取;
  • 在20层与第9层进行concat操作,再进行C2f模块特征提取。

3.目标检测:最后将15层、18层、21层分别经过Detect模块进行目标检测。

YOLO的模型构建过程解析

代码调用时序

构建模型的核心代码

通过查看上述YOLO构建模型过程,其核心代码主要是以下两处:
核心代码1:self._smart_load()

    def _new(self, cfg: str, task=None, model=None, verbose=False) -> None:
        # 以上部分省略...
        self.model = (model or self._smart_load("model"))(cfg_dict, verbose=verbose and RANK == -1)  # 核心代码
        self.overrides["model"] = self.cfg
        self.overrides["task"] = self.task

        self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}  # 核心代码
        # 以下部分省略...

代码解析:
self.model = (model or self._smart_load("model"))(cfg_dict, verbose=verbose and RANK == -1)

  • 特性:使用 Python 的条件表达式,选择不同的函数或对象进行赋值
  • 解释说明:
    • model or self._smart_load("model") 是一个条件表达式,它会根据 model 变量是否为真值来选择赋值的对象。
    • 如果 model 为真值(非空),则 self.model 将被赋值为 model
    • 如果 model 为假值(空),则会调用 self._smart_load("model") 方法来加载模型,并将返回的对象赋值给 self.model
    • 然后与,第二个括号(cfg_dict, verbose=verbose and RANK == -1) 拼接,形成self.model(cfg_dict, verbose=verbose and RANK == -1)

self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}

  • 特性:使用 Python 中的字典合并操作符 ** 来将两个字典合并。
  • 解释说明:
    • DEFAULT_CFG_DICTself.overrides 是两个字典。
    • **DEFAULT_CFG_DICTDEFAULT_CFG_DICT 字典中的所有键值对解包并添加到新的字典中。
    • **self.overrides 同样将 self.overrides 字典中的所有键值对解包并添加到同一个新的字典中。
    • 最终,{\*\*DEFAULT_CFG_DICT, \*\*self.overrides} 表示将这两个字典合并成一个新的字典,其中 self.overrides中的键值对将覆盖 DEFAULT_CFG_DICT 中的同名键值对。
  • 举例:
dict1 = {'a': 1, 'b': 2}
dict2 = {'b': 3, 'c': 4}
merged_dict = {**dict1, **dict2}
print(merged_dict)
# 运行结果:
# {'a': 1, 'b': 3, 'c': 4}

核心代码2:self.task_map[self.task][key]

def task_map(self):
        """Map head to model, trainer, validator, and predictor classes."""
        return {
            "classify": {
                "model": ClassificationModel,
                "trainer": yolo.classify.ClassificationTrainer,
                "validator": yolo.classify.ClassificationValidator,
                "predictor": yolo.classify.ClassificationPredictor,
            },
            "detect": {
                "model": DetectionModel,
                "trainer": yolo.detect.DetectionTrainer,
                "validator": yolo.detect.DetectionValidator,
                "predictor": yolo.detect.DetectionPredictor,
            },
            # 以下部分省略...
        }

代码解析:

  • 以上代码根据传入的self.task类型(例如:classify),来创建对应的模型、训练器、验证器、预测器类。
  • 特性:使用了 Python 中的字典,将字符串键与类对象值进行关联,以实现将类对象映射到不同的功能模块
  • 字典与类对象映射的举例
class Dog:
    def __init__(self, name):
        self.name = name
class Cat:
    def __init__(self, name):
        self.name = name
# 创建一个字典,将字符串键映射到不同的类对象
animal_map = {
    "dog": Dog,
    "cat": Cat,
}
# 根据键来实例化不同的类对象
my_dog = animal_map["dog"]("Buddy")
my_cat = animaljson_map["cat"]("Whiskers")
print(my_dog.name)  # 输出: Buddy
print(my_cat.name)  # 输出: Whiskers

核心代码3:parse_model()函数

def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)
    """Parse a YOLO model.yaml dictionary into a PyTorch model."""
    import ast

    # Args
    max_channels = float("inf")
    nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))
    depth, width, kpt_shape = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple", "kpt_shape"))
    if scales:
        scale = d.get("scale")
        if not scale:
            scale = tuple(scales.keys())[0]
            LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{scale}'.")
        depth, width, max_channels = scales[scale]

    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
        if verbose:
            LOGGER.info(f"{colorstr('activation:')} {act}")  # print

    ch = [ch]
    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
        m = getattr(torch.nn, m[3:]) if "nn." in m else globals()[m]  # get module
        for j, a in enumerate(args):
            if isinstance(a, str):
                with contextlib.suppress(ValueError):
                    args[j] = locals()[a] if a in locals() else ast.literal_eval(a)

        n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain
        if m in {
            Classify,
            Conv,
            ConvTranspose,
            GhostConv,
            Bottleneck,
            GhostBottleneck,
            SPP,
            SPPF,
            DWConv,
            Focus,
            BottleneckCSP,
            C1,
            C2,
            C2f,
            RepNCSPELAN4,
            ELAN1,
            ADown,
            AConv,
            SPPELAN,
            C2fAttn,
            C3,
            C3TR,
            C3Ghost,
            nn.ConvTranspose2d,
            DWConvTranspose2d,
            C3x,
            RepC3,
            PSA,
            SCDown,
            C2fCIB,
        }:
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            if m is C2fAttn:
                args[1] = make_divisible(min(args[1], max_channels // 2) * width, 8)  # embed channels
                args[2] = int(
                    max(round(min(args[2], max_channels // 2 // 32)) * width, 1) if args[2] > 1 else args[2]
                )  # num heads

            args = [c1, c2, *args[1:]]
        # ...((篇幅原因,代码已做省略))

        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
        t = str(m)[8:-2].replace("__main__.", "")  # module type
        m.np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type = i, f, t  # attach index, 'from' index, type
        if verbose:
            LOGGER.info(f"{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f}  {t:<45}{str(args):<30}")  # print
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)
    return nn.Sequential(*layers), sorted(save)

代码解析:

  • 参数设置:
    • 以上代码定义了一些初始参数,如 max_channels、nc(类别数)、act(激活函数类型)、scales(模型尺度)、depth(深度倍数)、width(宽度倍数)等。
    • 根据配置文件中的 scales 参数设置模型的深度、宽度和最大通道数。
  • 激活函数设置:
    • 如果配置文件中指定了激活函数类型 act,则重新定义默认激活函数为指定的激活函数(如 nn.SiLU())。
  • 模型解析:
    • 遍历配置文件中的 backbone 和 head 部分,解析每个层的信息。
    • 根据模块类型 m(如 Conv、Bottleneck 等)选择相应的处理逻辑,设置输入通道 c1、输出通道 c2 以及其他参数。
  • 创建模块:
    • 根据解析得到的参数和模块类型,通过nn.Sequential(*layers)创建相应的层,最终返回解析后的模型。

YOLO的模型训练过程解析

代码调用时序

构建模型的核心代码

核心代码1:ClassificationDataset数据集


    def __init__(self, root, args, augment=False, prefix=""):
        # 以上内容省略
        self.torch_transforms = (
            classify_augmentations(
                size=args.imgsz,
                scale=scale,
                hflip=args.fliplr,
                vflip=args.flipud,
                erasing=args.erasing,
                auto_augment=args.auto_augment,
                hsv_h=args.hsv_h,
                hsv_s=args.hsv_s,
                hsv_v=args.hsv_v,
            )
            if augment
            else classify_transforms(size=args.imgsz, crop_fraction=args.crop_fraction)
        )

代码解析:

  • 以上代码是ClassificationDataset数据集的初始化函数
  • self.torch_transforms = …:
    • 根据条件选择不同的数据增强方法:
    • 如果 augment 为 True,则调用 classify_augmentations() 方法进行数据增强。
    • 否则,调用 classify_transforms() 方法进行数据转换。
  • 除此之外,该数据集按照标准规范,实现了getitem(),len()等回调函数。

核心代码2:筹备训练

    def _setup_train(self, world_size):
        """Builds dataloaders and optimizer on correct rank process."""

        # Model
        self.run_callbacks("on_pretrain_routine_start")
        ckpt = self.setup_model()
        self.model = self.model.to(self.device)
        self.set_model_attributes()

        # Freeze layers
        # 篇幅原因,代码已省略

        # Check AMP
        # 篇幅原因,代码已省略

        # Check imgsz
        gs = max(int(self.model.stride.max() if hasattr(self.model, "stride") else 32), 32)  # grid size (max stride)
        self.args.imgsz = check_imgsz(self.args.imgsz, stride=gs, floor=gs, max_dim=1)
        self.stride = gs  # for multiscale training

        # Batch size
        # 篇幅原因,代码已省略

        # Dataloaders
        batch_size = self.batch_size // max(world_size, 1)
        self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode="train")
        # 代码已省略

        # Optimizer
        self.accumulate = max(round(self.args.nbs / self.batch_size), 1)  # accumulate loss before optimizing
        weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs  # scale weight_decay
        iterations = math.ceil(len(self.train_loader.dataset) / max(self.batch_size, self.args.nbs)) * self.epochs
        self.optimizer = self.build_optimizer(
            model=self.model,
            name=self.args.optimizer,
            lr=self.args.lr0,
            momentum=self.args.momentum,
            decay=weight_decay,
            iterations=iterations,
        )
        # Scheduler
        # 篇幅原因,代码已省略

核心代码3:开始训练

def _do_train(self, world_size=1):
        """Train completed, evaluate and plot if specified by arguments."""
        if world_size > 1:
            self._setup_ddp(world_size)
        self._setup_train(world_size)

        nb = len(self.train_loader)  # number of batches
        nw = max(round(self.args.warmup_epochs * nb), 100) if self.args.warmup_epochs > 0 else -1  # warmup iterations
        last_opt_step = -1
        self.epoch_time = None
        self.epoch_time_start = time.time()
        self.train_time_start = time.time()
        # 篇幅原因,代码已省略

        epoch = self.start_epoch
        self.optimizer.zero_grad()  # zero any resumed gradients to ensure stability on train start
        while True:
            self.epoch = epoch
            self.run_callbacks("on_train_epoch_start")
            # 篇幅原因,代码已省略

            self.model.train()
            if RANK != -1:
                self.train_loader.sampler.set_epoch(epoch)
            pbar = enumerate(self.train_loader)
            # Update dataloader attributes (optional)
            if epoch == (self.epochs - self.args.close_mosaic):
                self._close_dataloader_mosaic()
                self.train_loader.reset()

            # 篇幅原因,代码已省略
            self.tloss = None
            for i, batch in pbar:
                self.run_callbacks("on_train_batch_start")
                # Warmup
                # 篇幅原因,代码已省略

                # Forward 正向传播
                with torch.cuda.amp.autocast(self.amp):
                    batch = self.preprocess_batch(batch)
                    self.loss, self.loss_items = self.model(batch)  # 损失计算
                    if RANK != -1:
                        self.loss *= world_size
                    self.tloss = (
                        (self.tloss * i + self.loss_items) / (i + 1) if self.tloss is not None else self.loss_items
                    )

                # Backward 反向传播
                self.scaler.scale(self.loss).backward()

                # Optimize 优化一步
                if ni - last_opt_step >= self.accumulate:
                    self.optimizer_step()
                    last_opt_step = ni

                    # Timed stopping
                    # 篇幅原因,代码已省略

                # Log
                mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G"  # (GB)
                loss_len = self.tloss.shape[0] if len(self.tloss.shape) else 1
                losses = self.tloss if loss_len > 1 else torch.unsqueeze(self.tloss, 0)
                # 篇幅原因,代码已省略

            self.lr = {f"lr/pg{ir}": x["lr"] for ir, x in enumerate(self.optimizer.param_groups)}  # for loggers
            self.run_callbacks("on_train_epoch_end")
            if RANK in {-1, 0}:
                final_epoch = epoch + 1 >= self.epochs
                self.ema.update_attr(self.model, include=["yaml", "nc", "args", "names", "stride", "class_weights"])

                # Validation
                # 篇幅原因,代码已省略

                # Save model 保存模型
                if self.args.save or final_epoch:
                    self.save_model()
                    self.run_callbacks("on_model_save")

            # Scheduler
            # 篇幅原因,代码已省略

            # Early Stopping
            if RANK != -1:  # if DDP training
                broadcast_list = [self.stop if RANK == 0 else None]
                dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
                self.stop = broadcast_list[0]
            if self.stop:
                break  # must break all DDP ranks
            epoch += 1

        # 篇幅原因,代码已省略

代码解析:
在_do_train函数中,可以看到深度学习基本的步骤,即:

  1. 正向传播
  2. 损失计算
  3. 反向传播
  4. 优化一步
  5. 清空梯度
  6. 保存模型

    备注:清空梯度封装在optimizer_step函数中了。

内容小结

  • YOLOv8的网络结构
    • 主要有主干网络和Head网络组成
    • 主干网络中进行P1-P5层卷积,经过SPPF后进行Upsample上采样操作后,与P3、P4、P5进行concat操作
    • 最后通过15层、18层、21层分别经过Detect模块进行目标检测。
  • YOLOv8的代码解析
    • 构建模型的核心函数是:self._smart_load()、self.task_map[self.task][key]和parse_model函数
    • self.task_map使用了字符串键与类对象值进行关联,由此达到根据关键字选择对应的模型类
    • parse_model函数中通过读取参数、设置激活函数后,使用nn.Sequential创建对应的模型
    • 训练模型的核心函数为ClassificationDataset的封装、_setup_train()函数和_do_train()函数
    • _do_train()函数中包含深度学习的基本步骤,即:正向传播→损失计算→反向传播→优化一步→清空梯度→保存模型

参考资料

B站:YoloV8Ultralytics模型结构详细讲解

YOLOv8模型yaml结构图理解(逐层分析)

欢迎关注公众号以获得最新的文章和新闻

发表评论

您的邮箱地址不会被公开。 必填项已用 * 标注

分类文章

personal_logo
Dongming
自由职业者

推荐活动

推荐文章

【操作攻略】GPU云环境的使用介绍
【项目实战】通过LLaMaFactory+Qwen2-VL-2B微调一个多模态医疗大模型
【课程总结】day19(中):Transformer架构及注意力机制了解
【重拾数学知识】矢量的点乘和叉乘
【课程总结】Day11(上):手势图像识别实战(LeNet模型)
【课程总结】day34:多模态大模型之ViT模型、CLIP模型论文阅读理解
【目录】AI知识学习路线图
【课程总结】day33:文生图StableDiffusion模型初步了解以及部署体验
【课程总结】day32(下):Xinference部署向量化模型
【课程总结】day31:多模态大模型初步了解
内容目录
滚动至顶部