您的当前位置:首页正文

YOLOV5入门教学 - yolo.py文件

2024-11-10 来源:个人技术集锦

 一、导入包和本地配置

FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
if platform.system() != "Windows":
    ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

这段代码会获取当前文件的绝对路径,并使用Path库将其转换为Path对象。

这部分代码的主要作用有两个:

  • 将当前项目添加到系统路径上,以便项目中的模块可以调用。
  • 将当前项目的相对路径保存在ROOT中,便于寻找项目中的文件。
from models.common import (
    C3,
    C3SPP,
    C3TR,
    SPP,
    SPPF,
    Bottleneck,
    BottleneckCSP,
    C3Ghost,
    C3x,
    Classify,
    Concat,
    Contract,
    Conv,
    CrossConv,
    DetectMultiBackend,
    DWConv,
    DWConvTranspose2d,
    Expand,
    Focus,
    GhostBottleneck,
    GhostConv,
    Proto,
)
from models.experimental import MixConv2d
from utils.autoanchor import check_anchor_order
from utils.general import LOGGER, check_version, check_yaml, colorstr, make_divisible, print_args
from utils.plots import feature_visualization
from utils.torch_utils import (
    fuse_conv_and_bn,
    initialize_weights,
    model_info,
    profile,
    scale_img,
    select_device,
    time_sync,
)

try:
    import thop  # for FLOPs computation
except ImportError:
    thop = None

这些都是YOLOv5项目中自定义的库,由于前一步已经把路径加载上了,所以可以导入,这个顺序不可调整。具体来说,代码从以下几个文件中导入了部分函数和类:

  • models.common
    这个是YOLOv5的网络结构,包括 C3SPPFGhostConv 等模块,这些模块用于构建YOLOv5的神经网络。

  • models.experimental
    实验性的代码,包括 MixConv2d,这是一个组合了不同卷积核大小的卷积层,可能用于优化模型的性能。

  • utils.autoanchor
    定义了自动生成锚框的方法,比如 check_anchor_order 用于检查锚框的顺序是否符合要求。

  • utils.general
    定义了一些常用的工具函数,比如检查版本、是否在使用正确的配置文件、检查图像大小是否符合要求、打印命令行参数等。

  • utils.plots
    定义了 feature_visualization 函数,可以在训练过程中对特征图进行可视化,帮助理解模型的内部工作原理。

  • utils.torch_utils
    定义了一些与PyTorch有关的工具函数,如选择设备、同步时间、初始化权重、打印模型信息等。

通过引入这些模块,可以更方便地进行目标检测的相关任务,并且减少了代码的复杂度和冗余。

二、parse_model函数

def parse_model(d, ch):
    """Parses a YOLOv5 model from a dict `d`, configuring layers based on input channels `ch` and model architecture."""
    LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
    anchors, nc, gd, gw, act, ch_mul = (
        d["anchors"],
        d["nc"],
        d["depth_multiple"],
        d["width_multiple"],
        d.get("activation"),
        d.get("channel_multiple"),
    )
    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
        LOGGER.info(f"{colorstr('activation:')} {act}")  # print
    if not ch_mul:
        ch_mul = 8
    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
    no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)

? parse_model函数

parse_model 函数用于从模型配置字典(d)中解析YOLOv5的模型结构。通过提取yaml文件中的配置,将各个模块进行组装,最终构建出完整的模型网络架构。简单来说,就是把 yaml 文件中的网络结构实例化为对应的模型。

2.1 获取对应参数

def parse_model(d, ch):
    """Parses a YOLOv5 model from a dict `d`, configuring layers based on input channels `ch` and model architecture."""
    LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")
    anchors, nc, gd, gw, act, ch_mul = (
        d["anchors"],
        d["nc"],
        d["depth_multiple"],
        d["width_multiple"],
        d.get("activation"),
        d.get("channel_multiple"),
    )
    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
        LOGGER.info(f"{colorstr('activation:')} {act}")  # print
    if not ch_mul:
        ch_mul = 8
    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
    no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)

这段代码的主要作用是从配置字典 d 中提取模型的相关参数,并为后续的网络构建做好准备。具体步骤如下:

2.2 搭建网络前准备

    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 = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            with contextlib.suppress(NameError):
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings

这段代码主要是遍历backbonehead的每一层,获取搭建网络结构前的一系列信息。

我们还是先解释参数 layers, savec2

  • layers
    保存每一层的层结构。最终用于构建完整的模型。

  • save
    记录下所有层结构中 from 不等于 -1 的层结构序号。save 列表中的层会在训练过程中需要保存下来,用于后续层的拼接或其他操作。

  • c2
    记录当前层的输出通道数,用于下一层的输入通道数。

for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):
  • f:
    表示当前输入从哪些层获取,通常是一个层的序号。f可以是单个数字,表示输入来自哪一层的输出,也可以是列表,表示来自多层的输出。

  • n:
    表示当前层结构需要重复的次数。例如,在YOLOv5的C3模块中,n决定了重复多少次卷积操作。

  • m:
    当前层的模块类型,如 ConvFocus 等,这些模块会被动态加载。

  • args:
    表示当前层的参数列表,例如卷积层的输出通道数、卷积核大小等。


eval() 函数的使用

接着遇到一个函数 eval(),主要用于将字符串转换成有效的表达式来求值,并且返回执行的结果。

m = eval(m) if isinstance(m, str) else m  # eval strings
  • m = eval(m)
    • 如果 m 是字符串(即层的模块类型以字符串形式给出),通过 eval 将字符串转为实际的模块对象(如 ConvC3 等)。

同样地,args 中可能也包含字符串类型的参数,比如 'None', 'nc', 'anchors' 等。通过下面的代码段,将这些字符串转换为对应的值:

for j, a in enumerate(args):
    with contextlib.suppress(NameError):
        args[j] = eval(a) if isinstance(a, str) else a  # eval strings
  • args[j] = eval(a)
    args 列表中的每个参数 a,如果它是字符串类型,通过 eval() 进行转换。

  • contextlib.suppress(NameError)
    这里使用了 contextlib.suppress() 来忽略可能出现的 NameError。如果 eval() 无法正确执行字符串转换,则跳过该错误。

2.3 更新当前层的参数,计算c2

        n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
        if m in {
            Conv,
            GhostConv,
            Bottleneck,
            GhostBottleneck,
            SPP,
            SPPF,
            DWConv,
            MixConv2d,
            Focus,
            CrossConv,
            BottleneckCSP,
            C3,
            C3TR,
            C3SPP,
            C3Ghost,
            nn.ConvTranspose2d,
            DWConvTranspose2d,
            C3x,
        }:
            c1, c2 = ch[f], args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, ch_mul)

这段代码主要是更新当前层的 args,计算 c2(当前层的输出通道数):

  1. n = max(round(n * gd), 1)

    • n 是当前层的深度系数,通过乘以 gd(depth_multiple)来调整。gd 控制模型的深度,例如 YOLOv5s 的 gd 值为 0.33,意味着把模型的深度缩减为原来的 1/3。
    • round() 用来取整,并且至少为 1 层(避免 n=0 的情况)。
  2. 判断模块是否属于指定类型

    if m in {Conv, GhostConv, Bottleneck, GhostBottleneck, ... }:
    • 判断当前层 m 是否属于卷积类模块,如 ConvC3SPPFocus 等,这些模块通常会涉及输入通道和输出通道的变换。
  3. 提取当前层的输入通道数 c1 和输出通道数 c2

    c1, c2 = ch[f], args[0]
    • c1:当前层的输入通道数,ch[f] 表示从前面第 f 层输入。
    • c2:当前层的输出通道数,args[0] 是该层的输出通道数。
  4. 调整输出通道数 c2

    if c2 != no:  # if not output
        c2 = make_divisible(c2 * gw, ch_mul)
    • no:代表检测头的输出通道数,当 c2 不等于 no 时,需要调整。
    • 使用 make_divisible 函数调整通道数,确保 c2 是一个可被 ch_mul 整除的值。
    • c2 * gw 通过 gw(width_multiple)调整模型的宽度,类似于 gd 控制模型深度,gw 控制模型宽度,例如 YOLOv5s 的 gw 值为 0.5,意味着通道数会减半。

make_divisible 函数的作用

make_divisible(c2 * gw, ch_mul) 确保输出通道数 c2ch_mul 的倍数,使其适合硬件加速和模型优化。下面是 make_divisible 的实现:

def make_divisible(x, divisor):
    return math.ceil(x / divisor) * divisor
  • make_divisible
    • 作用是让 x 能够被 divisor 整除,通过向上取整,使 x 成为 divisor 的最小倍数。
    • 这种方式确保计算效率,并在某些硬件平台上优化性能。

为了更好地理解这段代码的实际运行过程,以下是一个简化的运行案例,展示如何动态调整YOLOv5模型的深度和宽度,并解释各个部分是如何工作的。

案例背景

假设我们正在构建YOLOv5的一个卷积层模块,例如 ConvC3,并且我们希望根据模型配置中的深度和宽度系数动态调整模型的深度和输出通道数。

输入数据

我们有一个 yaml 配置文件,其中包含如下模型参数:

depth_multiple: 0.33  # 控制深度
width_multiple: 0.5   # 控制宽度
anchors: [[10, 13, 16, 30], [30, 61, 62, 45], [116, 90, 156, 198]]  # 锚框
nc: 80  # 类别数

该配置适用于YOLOv5的 "small" 版本 (YOLOv5s),其中 depth_multiplewidth_multiple 控制模型的深度和宽度。

模拟模型层

假设我们正在处理一个 C3 模块,模块参数如下:

n = 3  # 层的重复次数,通常是一个初始值,表示该层中子层的个数
ch = [64, 128, 256]  # 通道数列表,表示前几层的输出通道数
args = [256]  # 目标输出通道数
f = 1  # 输入来自的层(假设来自前一层)

模型调整步骤

  1. 深度调整(n的变化)
gd = 0.33  # depth_multiple
n = max(round(n * gd), 1)  # 调整深度
  • n 初始值为 3,gd 是深度控制系数,为 0.33。
  • 计算出新的深度:n = max(round(3 * 0.33), 1),即 n = 1
    • 这意味着,原本应该重复 3 次的 C3 模块,现在只会重复 1 次,模型变得更浅,减少计算量。
  1. 检查是否属于调整类型的层
m = C3  # 当前的模块是C3
if m in {Conv, C3, ...}:
    c1, c2 = ch[f], args[0]  # c1是输入通道,c2是输出通道
  • 当前模块是 C3,满足判断条件。
  • c1(输入通道数)取自前一层的输出通道,即 ch[1] = 128
  • c2(输出通道数)从 args[0] 取值,即 c2 = 256
  1. 输出通道调整(c2的变化)
gw = 0.5  # width_multiple 控制模型宽度
c2 = make_divisible(c2 * gw, 8)  # 调整c2
  • 计算 c2c2 = 256 * 0.5 = 128
  • 然后通过 make_divisible 确保 c2 是 8 的倍数:
    c2 = make_divisible(128, 8)  # 结果仍然是 128
    • 这里没有变化,因为128已经是8的倍数。
  1. 结果输出

最终,我们调整后的层结构如下:

  • 输入通道数c1 = 128
  • 输出通道数c2 = 128
  • 深度:层重复 1 次,而不是 3 次。

这意味着,原本计划输出 256 个通道并重复 3 次的层,现在输出 128 个通道,并且只重复 1 次。通过这种方式,模型被动态缩小,适应资源受限的环境(例如移动设备)。

完整流程总结

  1. 深度调整:通过 depth_multiplegd),减少某些层的深度,降低计算开销。
  2. 宽度调整:通过 width_multiplegw),减少每层的输出通道数,使模型变窄,减少参数量和计算量。
  3. 输出调整:通过 make_divisible 确保输出通道数是硬件友好的倍数,便于优化计算效率。

这个案例展示了YOLOv5模型是如何根据配置动态调整深度和宽度的,使得同一个网络架构可以适应不同的计算资源条件。

2.4搭建该层网络

           args = [c1, c2, *args[1:]]
            if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
                args.insert(2, n)  # number of repeats
                n = 1
        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        # TODO: channel, gw, gd
        elif m in {Detect, Segment}:
            args.append([ch[x] for x in f])
            if isinstance(args[1], int):  # number of anchors
                args[1] = [list(range(args[1] * 2))] * len(f)
            if m is Segment:
                args[3] = make_divisible(args[3] * gw, ch_mul)
        elif m is Contract:
            c2 = ch[f] * args[0] ** 2
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        else:
            c2 = ch[f]

这段代码的主要功能是根据当前层的参数构建当前层的网络结构,并且为不同类型的模块(如 BottleneckCSPConcatDetect 等)动态计算输入和输出通道数。

args = [c1, c2, *args[1:]]
if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
    args.insert(2, n)  # number of repeats
    n = 1
  • 解释
    这里通过将当前层的输入通道 c1 和输出通道 c2 加入到 args 参数列表的前面。
    • 如果当前层属于 BottleneckCSPC3 等模块,表示这类模块具有重复结构,需要将重复次数 n 插入到 args 参数中,作为模块构造的参数。
    • 之后将 n 设置为 1,表示这些模块将不再进一步重复,因为重复次数已经交由该模块内部处理。
elif m is nn.BatchNorm2d:
    args = [ch[f]]
  • 解释
    如果当前层是 BatchNorm2d,则只需保留其输入通道数 ch[f],即从前一层输入的通道数。
elif m is Concat:
    c2 = sum(ch[x] for x in f)
  • 解释
    如果当前层是 Concat,表示该层用于将多个层的输出拼接起来。此时的输出通道数 c2 是所有输入通道数的和,即从 f 层输入的通道数之和。
elif m in {Detect, Segment}:
    args.append([ch[x] for x in f])
    if isinstance(args[1], int):  # number of anchors
        args[1] = [list(range(args[1] * 2))] * len(f)
    if m is Segment:
        args[3] = make_divisible(args[3] * gw, ch_mul)
  • 解释
    如果当前层是 Detect 或 Segment(检测或分割层),需要在 args 中添加从 f 层输入的通道数列表。
    • 对于 Detect 层,还会检查第二个参数 args[1] 是否为整数(表示锚框数量),如果是整数,则生成一个锚框索引列表,并重复多次,确保与输入的特征层数匹配。
    • 如果当前层是 Segment(语义分割),则对输出通道数进行处理,确保它是 ch_mul 倍数。
elif m is Contract:
    c2 = ch[f] * args[0] ** 2
  • 解释
    如果当前层是 Contract(收缩操作),其输出通道数 c2 会是输入通道数 ch[f] 乘以 args[0] 的平方。Contract 通常用于减少图像分辨率,同时增加通道数。
elif m is Expand:
    c2 = ch[f] // args[0] ** 2
  • 解释
    如果当前层是 Expand(扩展操作),其输出通道数 c2 会是输入通道数 ch[f] 除以 args[0] 的平方。Expand 用于增加图像分辨率,同时减少通道数。
else:
    c2 = ch[f]
  • 解释
    对于未列出的模块,默认情况下,输出通道数 c2 等于输入通道数 ch[f],即没有通道数变化。

代码逻辑总结

  • 处理模块的输入输出通道
    不同的模块需要不同的输入和输出通道处理方式,这段代码根据模块的类型动态调整其输入通道数 c1 和输出通道数 c2

  • 特定模块的逻辑

    • 对于卷积类模块(如 BottleneckCSPC3),输出通道数由配置中的 args[0] 提供,并根据宽度系数进行调整。
    • Concat 层的输出通道数是输入通道数的总和。
    • Detect 层用于目标检测,输出特征图包含多个 anchor(锚框),其通道数需要根据 anchors 和类别数进行调整。
    • Contract 和 Expand 是空间操作,它们调整图像尺寸和通道数。

应用案例

假设我们有以下模型配置:

  • C3 层:输入通道为 128,输出通道为 256,重复 3 次。
  • Concat 层:输入来自多个层,通道数分别为 64 和 128。
  • Detect 层:检测头,输出与锚框数量和类别数相关。

对于这些层,代码会根据每层的模块类型和输入输出通道信息,动态计算通道数并调整模型结构。

通过这种方式,YOLOv5 的模型构建能够灵活适应不同的网络配置,并确保通道数和计算效率的优化。

2.5输出模型layers

        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
        np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
        LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{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)

这段代码的主要功能是打印当前层结构的一些基本信息并将其保存。

构建的模块会被保存到 layers 列表中,并且将该层的输出通道数写入 ch 列表里。待全部循环结束后,返回一个包含所有层的模型。


详细解读

m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
  • 功能
    构建当前层的网络模块。
    • n > 1:如果当前层需要重复多次(例如 C3 模块),则通过 nn.Sequential 进行多次重复,并将每个重复的模块实例化并放入 nn.Sequential 中。
    • m(*args):如果不需要重复,则直接调用 m 模块,将参数 args 传入并实例化该模块。
t = str(m)[8:-2].replace("__main__.", "")  # module type
  • 功能
    通过将模块 m 转换为字符串来获取其类型。这里通过字符串切片和替换操作,提取模块的类型名称。
    • t 是模块的类型名称,如 ConvC3SPPF 等。
np = sum(x.numel() for x in m_.parameters())  # number params
  • 功能
    计算当前模块的参数总数。
    • m_.parameters():返回当前模块 m_ 的所有参数。
    • x.numel():获取每个参数张量的元素数量。
    • sum():将所有参数的元素数量加起来,得到总参数数目 np
m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
  • 功能
    给当前模块 m_ 附加一些额外的属性,用于记录模块的相关信息:
    • i:当前层的索引(层号)。
    • f:表示当前层的输入来源(上一层或多层)。
    • t:模块类型。
    • np:模块的参数数量。
LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}")  # print
  • 功能
    打印当前层的基本信息,包括:
    • i:当前层的索引号。
    • f:输入层的索引号。
    • n_:当前层的深度(即重复次数)。
    • np:参数数量。
    • t:模块类型。
    • args:模块的参数。

这一行的日志输出将层的构建细节打印出来,便于调试和查看模型结构。

save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
  • 功能
    将当前层的索引 f 保存到 save 列表中,但只保存那些 f != -1 的层(即非输入层的索引)。
    • x % i:这里确保保存的索引为合法层的索引号。
    • save:用于保存需要在后续网络构建中使用的层(主要是用于拼接等操作)。
layers.append(m_)
  • 功能
    将当前构建的模块 m_ 添加到 layers 列表中。layers 最终会被用于构建完整的 nn.Sequential 模型。
if i == 0:
    ch = []
ch.append(c2)
  • 功能
    维护输出通道数的列表 ch。如果当前是第一层 i == 0,则初始化 ch。每次处理完一层后,将该层的输出通道数 c2 添加到 ch 列表中。

最后返回构建好的模型

return nn.Sequential(*layers), sorted(save)
  • nn.Sequential(*layers)
    使用 nn.Sequential 将所有层(存储在 layers 列表中)组合成一个完整的模型。每一层都按顺序执行,作为整个模型的一部分。

  • sorted(save)
    返回所有需要保存的层的索引列表(保存自 save 列表),并对其进行排序,方便后续处理。返回的索引可以是用于拼接、跳接或其他网络结构中的特殊操作的层。


总结

  • 这段代码的主要作用是构建模型的每一层,并打印每一层的详细信息,如输入来源、模块类型、参数数量等。
  • 所有构建好的层会存储到 layers 列表中,最终组合成一个完整的 nn.Sequential 模型。
  • 还会记录哪些层需要保存,以便在后续网络结构中使用(如特征拼接或跳层连接)。

至此,整个模型的构建过程完成。

三、detect部分

Detect模块详解

Detect 模块是 YOLOv5 网络结构的最后一层,用于检测目标。它的功能是根据网络的输出特征图生成最终的预测框、置信度和类别。此模块通常在 yaml 配置文件的最后一行定义,格式为:

[*from], 1, Detect, [nc, anchors]

其中:

  • nc:表示分类数,即模型要检测的目标类别数量。
  • anchors:表示用于检测不同大小目标的锚框,通常由 yaml 文件中的配置决定。
class Detect(nn.Module):
    # YOLOv5 Detect head for detection models
    stride = None  # strides computed during build
    dynamic = False  # force grid reconstruction
    export = False  # export mode

    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
        """Initializes YOLOv5 detection layer with specified classes, anchors, channels, and inplace operations."""
        super().__init__()
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        self.na = len(anchors[0]) // 2  # number of anchors
        self.grid = [torch.empty(0) for _ in range(self.nl)]  # init grid
        self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]  # init anchor grid
        self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl,na,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
        self.inplace = inplace  # use inplace ops (e.g. slice assignment)

代码详解

1. 类变量的定义
stride = None  # strides computed during build
dynamic = False  # force grid reconstruction
export = False  # export mode
  • stride:指每个检测层的步幅,它会在模型构建过程中自动计算,代表特征图相对于输入图像的缩放比例。
  • dynamic:控制是否强制重构网格(grid),在某些情况下会被用来动态调整网格。
  • export:是否处于导出模式,导出模型时可能会进行一些优化。

2. 初始化函数 __init__
def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
  • nc:表示目标检测的类别数,默认值是 80(COCO 数据集包含 80 个类别)。
  • anchors:是锚框的列表,表示模型在不同尺度下用于检测的锚框大小。
  • ch:表示前面网络层的输出通道数,用于决定 Detect 模块的输入。
  • inplace:决定是否使用 inplace 操作(就地操作),默认是 True,可以优化内存使用。

3. 计算各个参数
self.nc = nc  # number of classes
self.no = nc + 5  # number of outputs per anchor
self.nl = len(anchors)  # number of detection layers
self.na = len(anchors[0]) // 2  # number of anchors
  • self.no:每个锚框的输出通道数。它等于类别数 nc 加上 5(x, y, w, h, confidence,即边界框的4个参数和置信度)。
  • self.nl:表示检测层的数量,通常是3个(对应3个不同的特征图)。
  • self.na:表示每个检测层上的锚框数量,通常为3个锚框。

4. 网格和锚框初始化
self.grid = [torch.empty(0) for _ in range(self.nl)]  # init grid
self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]  # init anchor grid
  • self.grid:初始化空的网格列表,用于存储每个检测层的网格坐标。网格的作用是将预测框映射回原始图像。
  • self.anchor_grid:初始化空的锚框网格,用于存储每个检测层的锚框坐标。

5. 注册锚框
self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))  # shape(nl, na, 2)
  • register_buffer:用来注册一个持久性缓冲区(buffer),它不会参与梯度计算,但会被保存和加载。这里将锚框信息注册为 anchors
  • torch.tensor(anchors).float().view(self.nl, -1, 2):将锚框数据转换为一个 tensor,并调整形状为 (nl, na, 2),表示每层的锚框数量和大小。

6. 卷积层的定义
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
  • self.m:使用 nn.ModuleList 来定义输出层,每个特征图都会对应一个卷积层。
    • 每个卷积层的输入通道是 ch 中定义的特征图通道数。
    • 输出通道数是 self.no * self.na,即每个锚框的输出通道数乘以锚框数量。

7. Inplace 操作
self.inplace = inplace  # use inplace ops (e.g. slice assignment)
  • inplace:表示是否使用就地操作(inplace),可以减少内存的占用。默认值为 True,在内存敏感的场景下非常有用。

参数解释

  • nc:分类数量(目标检测类别数量)。
  • no:每个锚框的输出数目,包括边界框的 4 个参数 (x, y, w, h)、置信度和分类数,即 no = 5 + nc
  • nl:检测层的数量(通常为 3),每一层对应不同的特征图。
  • na:每层的锚框数量,一般为 3。
  • grid:网格坐标,用来对应特征图上的每个预测框。
  • anchor_grid:锚框网格,用于表示每层特征图对应的锚框。

3.1 forward函数

forward 函数是 YOLOv5 中用于目标检测的核心计算部分。它接收网络输出的特征图,并根据不同的情况处理输入数据,生成最终的检测结果(边界框、置信度和类别)。该函数会根据是否处于训练模式推理模式,以不同的方式处理特征图。

    def forward(self, x):
        """Processes input through YOLOv5 layers, altering shape for detection: `x(bs, 3, ny, nx, 85)`."""
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                if isinstance(self, Segment):  # (boxes + masks)
                    xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
                    xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)
                else:  # Detect (boxes only)
                    xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
                    xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf), 4)
                z.append(y.view(bs, self.na * nx * ny, self.no))

        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)

1. 函数输入与输出

def forward(self, x):
    """Processes input through YOLOv5 layers, altering shape for detection: `x(bs, 3, ny, nx, 85)`."""
  • 输入

    • x 是从前面卷积层得到的特征图,包含3个不同尺度的特征图,形状可能是 (batch_size, 255, 80, 80)(batch_size, 255, 40, 40) 和 (batch_size, 255, 20, 20)。这里 255 是由 na * (nc + 5) 计算得来的,其中 na 是每个网格的锚框数量,nc 是目标类别数。
  • 输出

    • 如果处于训练模式,返回经过处理的 x
    • 如果处于推理模式(即测试或推理阶段),返回拼接后的检测结果 z,包含预测的边界框、置信度和类别信息。

2. 初始化推理输出列表

z = []  # inference output
  • z:用来存储推理输出结果的列表。最终所有检测结果会被存储在 z 中。

3. 遍历每个检测层的特征图

for i in range(self.nl):
    x[i] = self.m[i](x[i])  # conv
    bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
    x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
  • 卷积处理:每个特征图 x[i] 会经过一个卷积层 self.m[i],将其输出通道数调整为 (na * no),即 (锚框数 * 每个锚框的输出数)
  • 形状调整:调整后的特征图 x[i] 被 view() 转换成 [batch_size, 锚框数, 输出数, 高度, 宽度]。随后使用 permute() 将维度重新排序,以便按照 [batch_size, 锚框数, 高度, 宽度, 输出数] 的顺序进行操作。

4. 推理模式下的处理

if not self.training:  # inference
    if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
        self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
  • 动态网格处理:在推理模式下,检测模块会根据特征图的大小构建网格(grid),这个网格用于将模型预测的边界框坐标映射回原图。
    • 如果 self.dynamic 为 True 或者网格的大小不符合当前特征图,则调用 _make_grid() 函数来重新生成网格和锚框网格。

5. 位置和尺寸计算

检测(Detect)模式下的处理

xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
y = torch.cat((xy, wh, conf), 4)
  • sigmoid():对特征图中的 xy 坐标、wh 宽高、以及 conf 置信度部分分别应用 sigmoid 激活函数。
  • 位置计算 (xy):将 xy 预测值映射到网格坐标系中,并乘以步长 stride,将其转换为原图中的坐标。
  • 尺寸计算 (wh):对 wh 使用 sigmoid 进行缩放,然后通过平方操作得到预测框的宽高,并乘以锚框尺寸。
  • 拼接 (torch.cat):将处理好的 xywh 和 conf 拼接在一起,形成最终的预测结果。

6. 最终输出

z.append(y.view(bs, self.na * nx * ny, self.no))
  • 预测结果:将处理好的特征图 y 变形为 [batch_size, 锚框数 * 高度 * 宽度, 输出数],并添加到推理输出列表 z 中。

7. 返回结果

return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
  • 训练模式下:如果处于训练模式,函数会直接返回调整后的 x
  • 推理模式下
    • 如果是导出模式(self.export),返回拼接后的推理结果 z
    • 否则,返回拼接后的推理结果 z 和特征图 x

总结

  • 训练模式:特征图被处理后,直接返回,用于计算损失。
  • 推理模式:对特征图进行卷积、变换、坐标和尺寸计算后,生成最终的检测结果(包括预测框、置信度和类别),并返回。
  • 核心操作:包括对特征图的调整、动态构建网格、以及坐标和尺寸的映射。

3.2 make_grid

    def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
        """Generates a mesh grid for anchor boxes with optional compatibility for torch versions < 1.10."""
        d = self.anchors[i].device
        t = self.anchors[i].dtype
        shape = 1, self.na, ny, nx, 2  # grid shape
        y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
        yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)  # torch>=0.7 compatibility
        grid = torch.stack((xv, yv), 2).expand(shape) - 0.5  # add grid offset, i.e. y = 2.0 * x - 0.5
        anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
        return grid, anchor_grid

_make_grid 函数详解

_make_grid 函数用于生成 YOLOv5 检测层中的网格坐标和锚框网格。通过这个函数,特征图上的每个像素点可以映射到输入图像中的实际坐标,配合锚框进行目标检测。这段代码考虑了不同的 PyTorch 版本兼容性,并为特定的检测层生成相应的网格和锚框信息。


1. 函数输入与输出

def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
  • 输入参数

    • nx 和 ny:特征图的宽度和高度(例如,20x20 表示特征图的尺寸)。
    • i:当前处理的检测层索引,用于从 self.anchors 和 self.stride 中获取该层的信息。
    • torch_1_10:用于检测 PyTorch 版本,决定如何构建网格,是否兼容老版本。
  • 输出

    • 返回两个张量:
      • grid:表示特征图网格的坐标,形状为 (1, na, ny, nx, 2)
      • anchor_grid:表示锚框网格的坐标,形状为 (1, na, ny, nx, 2),是每个锚框的实际宽高。

2. 设备和数据类型

d = self.anchors[i].device
t = self.anchors[i].dtype
  • d:表示锚框的设备(即当前锚框存储在 GPU 还是 CPU 上)。
  • t:表示锚框的 dtype(数据类型,例如 float32)。

3. 定义网格的形状

shape = 1, self.na, ny, nx, 2  # grid shape
  • shape:定义了网格的形状,具体为:
    • 1:表示批次维度,通常是1。
    • self.na:表示锚框数量(例如,3个锚框)。
    • ny 和 nx:表示特征图的高度和宽度。
    • 2:表示网格坐标的两个维度 x 和 y

最终的网格形状为 (1, 锚框数, 高度, 宽度, 2)


4. 生成网格

y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)
  • torch.arange(ny) 和 torch.arange(nx):分别生成从 0 到 ny-1 和 nx-1 的坐标。
  • torch.meshgrid
    • 根据 PyTorch 版本,生成二维网格的坐标矩阵。网格的坐标表示每个像素点在特征图上的位置。
    • indexing="ij":确保输出的网格中 y 和 x 的排列方式是行优先的(行对应 y 轴,列对应 x 轴),这与传统的图像处理索引一致。

5. 构造网格坐标和锚框网格

grid = torch.stack((xv, yv), 2).expand(shape) - 0.5
  • torch.stack((xv, yv), 2):将 xv 和 yv 合并为一个 (ny, nx, 2) 形状的张量,其中每个元素表示网格点的 (x, y) 坐标。
  • expand(shape):将网格扩展到指定的形状 (1, na, ny, nx, 2),以适应所有锚框。
  • - 0.5:对网格坐标进行偏移操作,使得每个网格的中心点不再是整数,而是从 (-0.5, -0.5) 开始。这是 YOLO 检测框的一个常见操作,旨在对边界框进行调整,使其预测更精确。

6. 锚框网格

anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
  • self.anchors[i] * self.stride[i]:将锚框的宽高乘以当前层的步幅 stride,得到锚框的实际大小。
  • view((1, self.na, 1, 1, 2)):调整锚框的形状,使其与网格坐标匹配,形状为 (1, 锚框数, 1, 1, 2)
  • expand(shape):将锚框扩展到与网格相同的形状 (1, 锚框数, 高度, 宽度, 2)

7. 返回网格和锚框

return grid, anchor_grid
  • grid:返回的网格坐标,用于定位每个特征图上的像素点。
  • anchor_grid:返回的锚框网格,用于与网格坐标一起生成预测框。

总结

  • _make_grid 函数的核心功能是生成网格坐标 grid 和锚框网格 anchor_grid,并将它们用于 YOLO 的目标检测。
  • grid:表示特征图每个像素点在原始图像中的相对位置。
  • anchor_grid:表示锚框的实际大小,用于预测物体的边界框。
  • 通过这两者,YOLOv5 可以将特征图中的预测转换为输入图像中的实际边界框。

四、主要model介绍

这几个模型类(BaseModelDetectionModelSegmentationModelClassificationModel)是 YOLOv5 框架中的核心组件,分别用于处理不同的任务:目标检测、语义分割和图像分类。每个模型类有其特定的用途,能够解决不同的计算机视觉任务。下面是每个模型的具体用途和功能:


1. BaseModel

用途:这是 YOLOv5 中所有模型的基类,定义了一些通用的功能,比如前向传播、层的性能分析、模型信息打印、层融合等。BaseModel 为其他模型(如 DetectionModelSegmentationModelClassificationModel)提供了核心的基础功能。

  • 典型用途
    • 为 YOLOv5 中的其他模型类提供统一的接口和功能,如前向传播(forward)、性能分析(_profile_one_layer)、层融合(fuse)。
    • 通过继承 BaseModel,其他模型类能够复用这些基础功能,减少代码冗余。

2. DetectionModel

用途DetectionModel 是 YOLOv5 的核心模型类,用于目标检测任务。目标检测的任务是识别图像中的多个目标,并为每个目标预测一个边界框、置信度以及所属类别。

  • 功能

    • 目标检测:使用 YOLOv5 的主干(Backbone)提取特征,并通过检测头预测目标的位置(边界框)和类别。
    • 推理模式和增强推理:支持标准推理模式和增强推理模式,增强推理通过多尺度、图像翻转等手段提高检测精度。
    • 锚框处理:负责处理锚框,调整每个尺度上的锚框大小,适应不同图像中的目标。
    • 偏置初始化:初始化检测头的偏置参数,帮助模型更好地学习目标位置和分类信息。
  • 典型用途

    • 用于各类目标检测任务,比如自动驾驶中的行人检测、车辆检测,安防中的人物监控等。
    • 在 COCO、Pascal VOC 等标准数据集上进行目标检测的训练和推理。

3. SegmentationModel

用途SegmentationModel 是用于语义分割任务的模型。语义分割的任务是为每个像素分类,即为每个像素分配一个类别标签,通常用于精细化目标区域的标注。

  • 功能

    • 语义分割:基于 YOLOv5 的结构,增加了分割头,用于生成像素级别的分类结果。
    • 检测与分割结合:它可以同时执行目标检测和语义分割,支持多任务学习,能够预测边界框和分割掩码。
  • 典型用途

    • 在自动驾驶中用于道路分割、车道标识检测等任务。
    • 医学图像中用于识别特定区域,如病变区域、器官分割等。
    • 卫星图像处理中的地物分类,如道路、水体、建筑物的分割。

4. ClassificationModel

用途ClassificationModel 是用于图像分类任务的模型。图像分类的任务是给整个图像分配一个类别标签,通常用于识别图像的主体内容。

  • 功能

    • 图像分类:通过 YOLOv5 的 Backbone(特征提取部分)提取图像特征,并将其输入到分类头进行图像分类。
    • 模型转换:支持从 YOLOv5 的目标检测模型转换为分类模型,即只使用 Backbone 部分,省去检测头。
  • 典型用途

    • 用于经典的图像分类任务,如 ImageNet 数据集上的分类任务,识别图像中的物体种类。
    • 在工业生产中用于产品质量检测,通过分类模型识别产品是否符合标准。
    • 在安防监控中用于人脸分类或姿态识别。

总结

  • BaseModel:定义了 YOLOv5 中所有模型的核心功能,如前向传播、层融合、性能分析等。为其他模型类提供基础功能。
  • DetectionModel:用于目标检测任务,能够识别图像中的多个目标,并为每个目标预测边界框、置信度和类别标签。常用于物体检测和定位任务。
  • SegmentationModel:用于语义分割任务,能够为每个像素分配类别标签,常用于图像中目标区域的精细化分割。
  • ClassificationModel:用于图像分类任务,能够对整个图像进行分类,常用于识别图像的主要内容。

实际应用场景

  • 目标检测:应用于自动驾驶、安防监控、无人机视觉等场景,用来检测图像中的特定目标(如行人、车辆等)。
  • 语义分割:在医学图像、自动驾驶、遥感图像处理中,用于细致的目标分割,区分图像中的不同区域。
  • 图像分类:广泛应用于图像识别、产品质量检测、安防监控等场景,用来判断图像的主要内容。

4.1 BaseModel

class BaseModel(nn.Module):
    """YOLOv5 base model."""

    def forward(self, x, profile=False, visualize=False):
        """Executes a single-scale inference or training pass on the YOLOv5 base model, with options for profiling and
        visualization.
        """
        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    def _forward_once(self, x, profile=False, visualize=False):
        """Performs a forward pass on the YOLOv5 model, enabling profiling and feature visualization options."""
        y, dt = [], []  # outputs
        for m in self.model:
            if m.f != -1:  # if not from previous layer
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
            if profile:
                self._profile_one_layer(m, x, dt)
            x = m(x)  # run
            y.append(x if m.i in self.save else None)  # save output
            if visualize:
                feature_visualization(x, m.type, m.i, save_dir=visualize)
        return x

    def _profile_one_layer(self, m, x, dt):
        """Profiles a single layer's performance by computing GFLOPs, execution time, and parameters."""
        c = m == self.model[-1]  # is final layer, copy input as inplace fix
        o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0  # FLOPs
        t = time_sync()
        for _ in range(10):
            m(x.copy() if c else x)
        dt.append((time_sync() - t) * 100)
        if m == self.model[0]:
            LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  module")
        LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}")
        if c:
            LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")

    def fuse(self):
        """Fuses Conv2d() and BatchNorm2d() layers in the model to improve inference speed."""
        LOGGER.info("Fusing layers... ")
        for m in self.model.modules():
            if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                delattr(m, "bn")  # remove batchnorm
                m.forward = m.forward_fuse  # update forward
        self.info()
        return self

    def info(self, verbose=False, img_size=640):
        """Prints model information given verbosity and image size, e.g., `info(verbose=True, img_size=640)`."""
        model_info(self, verbose, img_size)

    def _apply(self, fn):
        """Applies transformations like to(), cpu(), cuda(), half() to model tensors excluding parameters or registered
        buffers.
        """
        self = super()._apply(fn)
        m = self.model[-1]  # Detect()
        if isinstance(m, (Detect, Segment)):
            m.stride = fn(m.stride)
            m.grid = list(map(fn, m.grid))
            if isinstance(m.anchor_grid, list):
                m.anchor_grid = list(map(fn, m.anchor_grid))
        return self

BaseModel 类详解

BaseModel 是 YOLOv5 中所有模型的基础类。它封装了模型的前向传播、层融合、模型信息打印、性能分析等常见功能。这个类为目标检测、分割、分类等不同任务的模型提供了通用的逻辑支持。

下面是对各个函数的详细解释:


1. forward()

def forward(self, x, profile=False, visualize=False):
    """
    执行单尺度推理或训练阶段的前向传播。
    可以选择进行性能分析(profile)或特征可视化(visualize)。
    """
    return self._forward_once(x, profile, visualize)  # 单尺度推理或训练
  • 功能:该函数是模型前向传播的入口。无论是训练阶段还是推理阶段,都会调用该函数进行单尺度的前向传播。
  • 参数
    • x:输入张量,通常是图像或特征图。
    • profile:如果为 True,会进行性能分析,记录 FLOPs 和执行时间。
    • visualize:如果为 True,会进行特征可视化,保存中间层的特征图。
  • 返回值:调用 _forward_once(),返回经过模型处理后的输出。

2. _forward_once()

def _forward_once(self, x, profile=False, visualize=False):
    """
    执行一次完整的前向传播,包含性能分析和特征可视化的选项。
    """
    y, dt = [], []  # 保存输出和执行时间
    for m in self.model:  # 遍历模型中的每一层
        if m.f != -1:  # 如果不是从前一层获取输入
            x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # 从之前的层获取输入
        if profile:
            self._profile_one_layer(m, x, dt)  # 进行性能分析
        x = m(x)  # 执行当前层
        y.append(x if m.i in self.save else None)  # 保存中间层输出
        if visualize:
            feature_visualization(x, m.type, m.i, save_dir=visualize)  # 可视化特征
    return x  # 返回最终输出
  • 功能:负责执行一次完整的前向传播,模型会按顺序依次通过每一层,生成最终的输出。
  • 流程
    • 对每一层 m 执行操作。如果该层的输入不来自前一层(m.f != -1),则根据 m.f 获取对应层的输出作为输入。
    • 如果 profile 为 True,则调用 _profile_one_layer 函数分析该层的性能。
    • 运行当前层 m(x),并将输出保存到列表 y 中。
    • 如果 visualize 为 True,则将该层的特征图可视化。
  • 返回值:返回经过所有层处理后的最终输出 x

3. _profile_one_layer()

def _profile_one_layer(self, m, x, dt):
    """
    分析单层的性能,计算 FLOPs、执行时间和参数数量。
    """
    c = m == self.model[-1]  # 如果是最后一层,复制输入
    o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0  # 计算FLOPs
    t = time_sync()
    for _ in range(10):
        m(x.copy() if c else x)  # 运行10次
    dt.append((time_sync() - t) * 100)  # 记录时间
    if m == self.model[0]:
        LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  module")  # 打印表头
    LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}")  # 打印性能信息
    if c:
        LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")  # 如果是最后一层,打印总时间
  • 功能:分析单层的性能,计算 FLOPs(浮点运算次数)、执行时间和参数数量。对于每一层都会记录这些信息,便于优化和调试。
  • 流程
    1. 使用 thop 库计算 FLOPs(每秒浮点运算次数)。
    2. 使用 time_sync() 记录每层的执行时间,循环运行 10 次,取平均值。
    3. 打印每一层的参数数量、FLOPs 和时间信息。

4. fuse()

def fuse(self):
    """
    融合卷积层和批量归一化层,以提高推理速度。
    """
    LOGGER.info("Fusing layers... ")
    for m in self.model.modules():
        if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"):  # 查找卷积层和批量归一化层
            m.conv = fuse_conv_and_bn(m.conv, m.bn)  # 融合卷积和BN层
            delattr(m, "bn")  # 删除BN层
            m.forward = m.forward_fuse  # 更新前向传播函数
    self.info()  # 打印模型信息
    return self
  • 功能:将卷积层(Conv2d)和批量归一化层(BatchNorm2d)进行融合,以提高推理速度。通过融合后,推理时不再需要额外的批量归一化操作,从而提升性能。
  • 操作流程
    • 遍历模型中的每一层,找到卷积层和批量归一化层的组合。
    • 通过 fuse_conv_and_bn 函数将两者融合成一个卷积层,并删除批量归一化层。
    • 更新前向传播函数为 forward_fuse,以便使用融合后的层进行推理。

5. info()

def info(self, verbose=False, img_size=640):
    """
    打印模型信息,包含模型的层数、参数量、计算量等。
    """
    model_info(self, verbose, img_size)
  • 功能:打印模型的详细信息,包括每一层的类型、参数数量、输出大小、FLOPs 等。用于了解模型的结构和性能。
  • 参数
    • verbose:如果为 True,则打印更加详细的模型信息。
    • img_size:用于计算特征图大小的输入图像尺寸,默认是 640。

6. _apply()

def _apply(self, fn):
    """
    对模型的张量应用指定的变换函数,比如转换到GPU(cuda)、CPU(cpu)或者半精度(half)计算。
    """
    self = super()._apply(fn)
    m = self.model[-1]  # Detect层
    if isinstance(m, (Detect, Segment)):
        m.stride = fn(m.stride)  # 应用变换到stride
        m.grid = list(map(fn, m.grid))  # 应用变换到grid
        if isinstance(m.anchor_grid, list):
            m.anchor_grid = list(map(fn, m.anchor_grid))  # 应用变换到anchor grid
    return self
  • 功能:应用函数 fn 对模型中的张量进行操作,比如将模型从 CPU 转移到 GPU,或将数据类型转换为半精度(float16)。它继承自 PyTorch 的 nn.Module 的 _apply() 方法,并对 YOLOv5 特定的 Detect 层和 Segment 层进行了额外处理。
  • 核心逻辑
    • 首先调用父类的 _apply() 方法,将变换应用到模型中的所有张量。
    • 如果模型包含 Detect 或 Segment 层,会进一步对 stridegrid 和 anchor_grid 进行变换处理。

总结

  • BaseModel 提供了 YOLOv5 模型的核心功能,如前向传播、性能分析、层融合和模型信息打印等。这些功能为具体任务(如目标检测、分割、分类)的模型类提供了基础设施。
  • 前向传播_forward_once() 是模型核心的前向传播函数,它遍历模型的每一层,并支持性能分析和特征可视化。
  • 性能分析:通过 _profile_one_layer() 函数,可以对每层的执行时间

4.2 DetectionModel

class DetectionModel(BaseModel):
    # YOLOv5 detection model
    def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None):
        """Initializes YOLOv5 model with configuration file, input channels, number of classes, and custom anchors."""
        super().__init__()
        if isinstance(cfg, dict):
            self.yaml = cfg  # model dict
        else:  # is *.yaml
            import yaml  # for torch hub

            self.yaml_file = Path(cfg).name
            with open(cfg, encoding="ascii", errors="ignore") as f:
                self.yaml = yaml.safe_load(f)  # model dict

这段代码是 YOLOv5 中 DetectionModel 类的初始化函数 __init__,它的主要功能是从 yaml 文件或字典中加载 YOLOv5 模型的配置,并根据这些配置构建模型。下面是对代码的详细解释:


代码功能概述

__init__ 函数的主要功能是为 YOLOv5 模型初始化一些关键参数,包括网络结构配置、输入通道数、分类数、锚框(anchors)和特征图步幅(stride)等。


1. 函数参数说明

def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None):
  • cfg:YOLOv5 模型配置文件,默认为 "yolov5s.yaml",这是 YOLOv5 最小版本的模型。可以是一个 YAML 文件路径,或者是已经解析的字典。
  • ch:输入图像的通道数,默认为 3(RGB 图像)。
  • nc:类别数量,表示数据集中的类别数。默认情况下,这个值会从 yaml 配置文件中读取。如果用户提供自定义的 nc,会覆盖配置文件中的默认值。
  • anchors:锚框的定义,用于目标检测任务中的不同大小目标。如果用户提供了自定义锚框,则会覆盖配置文件中的默认锚框。

2. 模型参数的加载

super().__init__()
  • 功能:调用父类 BaseModel 的初始化函数,初始化一些基础的模型功能。

3. 加载 yaml 文件或字典

if isinstance(cfg, dict):
    self.yaml = cfg  # 如果 cfg 是字典,直接使用
else:  # 如果 cfg 是 yaml 文件路径
    import yaml  # for torch hub

    self.yaml_file = Path(cfg).name
    with open(cfg, encoding="ascii", errors="ignore") as f:
        self.yaml = yaml.safe_load(f)  # 加载 yaml 文件内容并转为字典
  • 判断 cfg 类型

    • 如果 cfg 是一个字典,说明配置已经被解析,直接使用字典。
    • 如果 cfg 是一个 YAML 文件路径,则需要先导入 yaml 模块,然后从文件中加载 YAML 格式的模型配置,转换为字典格式。
  • self.yaml:存储了模型的配置参数,后续模型的构建和初始化将依据该字典中的参数进行。


4. 小结

  • 该 __init__ 函数的主要作用是从 yaml 文件或字典中读取 YOLOv5 的模型配置,接着将这些配置用于构建和初始化目标检测模型的各个部分。
  • 它能够灵活地支持两种输入:已经解析好的字典格式或未解析的 YAML 文件。通过这种设计,用户可以方便地根据自己的需求来加载和配置模型。

4.2.1搭建模型

        # Define model
        ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
        if nc and nc != self.yaml["nc"]:
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
            self.yaml["nc"] = nc  # override yaml value
        if anchors:
            LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}")
            self.yaml["anchors"] = round(anchors)  # override yaml value
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
        self.names = [str(i) for i in range(self.yaml["nc"])]  # default names
        self.inplace = self.yaml.get("inplace", True)

这段代码的功能

这段代码在 YOLOv5 的 DetectionModel 类初始化过程中,定义并构建模型结构。它从 yaml 配置文件或字典中读取模型的配置信息,处理输入通道数、类别数和锚框等参数,然后通过调用 parse_model 来构建模型的各层。


1. 定义模型输入通道数

ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
  • 功能:从 yaml 配置中获取输入通道数。如果 yaml 中没有定义输入通道数,则使用传入的默认值 ch(通常为 3,表示 RGB 图像)。
  • 作用:确保模型的输入通道数被正确设置,适应不同类型的输入图像。

2. 检查并覆盖类别数(nc

if nc and nc != self.yaml["nc"]:
    LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
    self.yaml["nc"] = nc  # override yaml value
  • 功能:检查用户是否在函数参数中传入了 nc(类别数)。如果传入了 nc 且与 yaml 文件中的类别数不一致,则使用传入的 nc 覆盖 yaml 中的值。
  • 作用:灵活处理用户自定义类别数,确保模型在处理不同数据集时能够设置正确的类别数量。

3. 检查并覆盖锚框(anchors

if anchors:
    LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}")
    self.yaml["anchors"] = round(anchors)  # override yaml value
  • 功能:检查用户是否提供了自定义的锚框。如果提供了,使用传入的 anchors 覆盖 yaml 文件中的锚框设置。
  • 作用:允许用户为不同的检测任务指定不同的锚框大小,以提高检测精度。

4. 解析模型结构

self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
  • 功能:调用 parse_model 函数,基于 yaml 文件中的配置构建模型结构,并返回模型和需要保存的层列表。
    • deepcopy(self.yaml):深度复制 yaml 文件中的配置,防止原始配置被修改。
    • ch=[ch]:传递输入通道数信息,作为模型构建的基础。
  • 作用:根据配置文件或字典,动态构建 YOLOv5 模型的每一层。

5. 设置类别名称

self.names = [str(i) for i in range(self.yaml["nc"])]  # default names
  • 功能:为每个类别生成一个默认名称,使用数字表示类别。例如,对于 80 个类别的 COCO 数据集,类别名称将为 '0', '1', '2', ..., '79'
  • 作用:为后续的可视化或结果输出提供类别标签,方便识别预测结果。

6. 处理 inplace 参数

self.inplace = self.yaml.get("inplace", True)
  • 功能:从 yaml 配置文件中获取 inplace 参数,用于控制是否使用原地操作(inplace operations)。如果 yaml 中没有定义 inplace,则默认为 True
  • 作用:在原地操作时可以减少内存开销,提高运行效率。

小结

  • 该代码首先检查并处理模型的输入通道数、类别数和锚框等配置,确保模型能够根据用户需求动态调整。
  • 然后,它通过 parse_model 函数构建 YOLOv5 的实际网络结构,并生成类别标签。
  • 这些设置为后续的推理和训练提供了基础,保证了模型的灵活性和可配置性。

4.2.2计算缩放倍数

        # Build strides, anchors
        m = self.model[-1]  # Detect()
        if isinstance(m, (Detect, Segment)):

            def _forward(x):
                """Passes the input 'x' through the model and returns the processed output."""
                return self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)

            s = 256  # 2x min stride
            m.inplace = self.inplace
            m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))])  # forward
            check_anchor_order(m)
            m.anchors /= m.stride.view(-1, 1, 1)
            self.stride = m.stride
            self._initialize_biases()  # only run once
                # Init weights, biases
        initialize_weights(self)
        self.info()
        LOGGER.info("")

这段代码的主要功能是构建模型的特征图步幅(stride)和锚框(anchor)信息,确保 YOLOv5 模型在推理和训练时能够正确计算边界框的位置和大小。接下来,我将逐步解析这些代码的实现及其意义。


主要步骤:

  1. 获取模型的最后一层 Detect

    m = self.model[-1]  # Detect()
    • self.model[-1]:表示获取模型的最后一层(一般是 Detect 层),这层用于处理目标检测的最终输出。
    • Detect 层:YOLOv5 用于预测目标边界框、置信度和类别的核心层。

  1. 定义一个 256×256 大小的输入

    s = 256  # 2x min stride
    • 定义 s = 256:表示模型的输入大小为 256×256,这个大小用于模拟推理时的输入图像尺寸。
    • 为什么使用 256:这是为了计算每一层输出特征图相对于输入图像的缩放比例,即步幅(stride)。通过定义一个固定大小的输入,模型可以推理出每个检测层的特征图步幅。

  1. 通过一次前向传播计算步幅

    def _forward(x):
        return self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)
    
    m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))])  # forward
    • _forward(x) 函数:负责执行前向传播。
      • 如果模型的最后一层是 Segment(用于语义分割),则只返回分割部分的输出;否则返回整个模型的输出。
    • torch.zeros(1, ch, s, s):创建了一个大小为 [1, ch, 256, 256] 的全零输入,ch 表示输入图像的通道数,通常为 3(RGB 图像)。
    • 步幅计算:通过 x.shape[-2] 获取前向传播后的特征图的高度。利用输入图像的高度 256 除以特征图的高度,得出每一层特征图相对于原始图像的缩放倍数(步幅)。
    • m.stride:步幅张量,存储了每层特征图的缩放倍数(步幅)。

  1. 调整锚框大小

    check_anchor_order(m)
    m.anchors /= m.stride.view(-1, 1, 1)
    self.stride = m.stride
    • check_anchor_order(m):检查锚框的顺序是否正确,确保大锚框用于较大的特征图,小锚框用于较小的特征图。
    • 调整锚框大小
      • m.anchors /= m.stride.view(-1, 1, 1):将每层的锚框除以该层的步幅。这是因为在特征图上,锚框的大小是相对于特征图的分辨率而言的,因此需要用步幅将锚框的大小调整到相应的尺度。
      • view(-1, 1, 1):将 stride 张量调整为适合和 anchors 张量相除的形状。
    • 存储步幅:将计算好的步幅 m.stride 存储到模型的属性 self.stride 中,供后续使用。

  1. 初始化偏置

    self._initialize_biases()  # only run once
    • 功能:初始化 Detect 层的偏置项(bias),这个步骤仅会在模型构建时运行一次。
    • 作用:通过初始化偏置项,使模型在训练时能够更快速地收敛,特别是针对目标检测任务的边界框预测。

代码的整体流程总结

  1. 获取最后一层的 Detect 层:用于目标检测的输出层。
  2. 定义 256×256 大小的输入:作为示例输入,模拟真实图像的推理过程。
  3. 通过一次前向传播计算特征图的步幅:利用输入图像大小和输出特征图大小的比例,计算每层特征图的缩放倍数(stride)。
  4. 调整锚框大小:根据每层特征图的步幅调整锚框的大小,使得锚框能适应不同尺度的特征图。
  5. 初始化偏置:确保偏置的合理初始化,提高模型的训练效果。

4.3数据增强

    def forward(self, x, augment=False, profile=False, visualize=False):
        """Performs single-scale or augmented inference and may include profiling or visualization."""
        if augment:
            return self._forward_augment(x)  # augmented inference, None
        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    def _forward_augment(self, x):
        """Performs augmented inference across different scales and flips, returning combined detections."""
        img_size = x.shape[-2:]  # height, width
        s = [1, 0.83, 0.67]  # scales
        f = [None, 3, None]  # flips (2-ud, 3-lr)
        y = []  # outputs
        for si, fi in zip(s, f):
            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
            yi = self._forward_once(xi)[0]  # forward
            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
            yi = self._descale_pred(yi, fi, si, img_size)
            y.append(yi)
        y = self._clip_augmented(y)  # clip augmented tails
        return torch.cat(y, 1), None  # augmented inference, train

代码功能解析

这段代码主要实现了 YOLOv5 模型在推理(Inference)阶段的两种模式:单尺度推理增强推理。其中,增强推理(Test Time Augmentation, TTA)通过对输入图像进行不同尺度的缩放和翻转,结合多次推理结果,提高模型在推理阶段的准确性。


1. forward() 函数

def forward(self, x, augment=False, profile=False, visualize=False):
    """执行单尺度或增强推理,并支持性能分析或特征可视化。"""
    if augment:
        return self._forward_augment(x)  # 使用增强推理
    return self._forward_once(x, profile, visualize)  # 单尺度推理
  • 功能:根据参数 augment 的值,决定是执行单尺度推理(_forward_once)还是增强推理(_forward_augment)。
  • 参数
    • x:输入的图像张量,通常是形状为 [batch_size, channels, height, width] 的 4D 张量。
    • augment:是否启用增强推理。为 True 时,启用增强推理,否则执行常规的单尺度推理。
    • profile:是否进行性能分析。为 True 时,记录每一层的性能(例如 FLOPs 和时间)。
    • visualize:是否进行特征可视化。为 True 时,会输出每一层的特征图。
  • 返回值:推理结果。如果是增强推理,返回多次推理结果的合并;如果是单尺度推理,返回单次推理结果。

2. _forward_augment() 函数

def _forward_augment(self, x):
    """执行增强推理,包括不同尺度和翻转的推理,返回合并的检测结果。"""
    img_size = x.shape[-2:]  # 输入图像的高和宽
    s = [1, 0.83, 0.67]  # 尺度列表:原始尺度、0.83 和 0.67 倍缩放
    f = [None, 3, None]  # 翻转:None 表示不翻转,3 表示左右翻转
    y = []  # 存储输出的列表
    for si, fi in zip(s, f):
        xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))  # 尺度缩放和翻转
        yi = self._forward_once(xi)[0]  # 单次推理
        yi = self._descale_pred(yi, fi, si, img_size)  # 反缩放并还原预测
        y.append(yi)
    y = self._clip_augmented(y)  # 裁剪增强推理结果的尾部
    return torch.cat(y, 1), None  # 返回合并的推理结果
  • 功能:该函数实现了增强推理,即通过对输入图像进行不同尺度的缩放和翻转,进行多次推理,然后将这些结果合并,返回最终的检测结果。
  • 核心步骤
    1. 尺度和翻转设置:定义了三个尺度 [1, 0.83, 0.67],分别表示原始尺寸、83% 缩小和67% 缩小,同时设置了翻转选项 [None, 3, None],其中 3 代表左右翻转。
    2. 循环处理:遍历每个尺度和翻转组合,先对输入图像进行缩放和翻转,然后调用 _forward_once 进行推理。
    3. 反缩放与还原预测:通过 _descale_pred 函数将预测结果还原回原图尺度。
    4. 结果合并:将不同尺度的推理结果存储在列表 y 中,最后使用 torch.cat 将它们沿着通道维度进行拼接。
    5. 裁剪:通过 _clip_augmented 函数对增强推理的结果进行裁剪,去除不必要的尾部。

增强推理的详细步骤

增强推理(TTA,Test Time Augmentation)是为了在推理阶段通过多种数据变换提高模型的鲁棒性和检测准确率。具体来说:

  1. 输入图像变换

    • 图像在原始尺度的基础上进行不同的缩放:1 表示原始大小,0.83 表示缩小至原尺寸的83%,0.67 表示缩小至原尺寸的67%。
    • 同时使用翻转:None 表示不翻转,3 表示水平翻转(左右翻转)。
  2. 多次推理

    • 对每次变换后的图像进行单次推理(通过 _forward_once 函数)。
    • 推理的结果为检测到的边界框、置信度和类别标签。
  3. 还原预测结果

    • 每次推理的结果都需要进行反缩放和翻转还原,使得这些结果能够和原始图像对齐。
    • 通过 _descale_pred 函数对预测结果进行还原。
  4. 合并结果

    • 将不同变换下的推理结果拼接在一起,通过 torch.cat(y, 1) 将所有推理结果在通道维度上进行合并。
  5. 裁剪

    • 为了去除多余的边界框预测结果,使用 _clip_augmented 函数对增强推理的结果进行裁剪,保留有效的部分。

代码关键点

  • 尺度缩放和翻转

    • 图像会根据预先定义的缩放比例 s = [1, 0.83, 0.67] 进行调整,并且在部分情况下进行左右翻转。这些操作增加了模型的鲁棒性,能够适应不同尺度和方向的物体。
  • 反缩放还原

    • 每次推理后,预测结果需要被反向缩放到原始输入图像的尺寸,这样才能保持结果的正确性。_descale_pred 负责这一过程。
  • 结果合并

    • 通过将所有增强推理的结果合并到一起,能够获取多视角的预测,提高检测的准确性。

小结

  • 该部分代码通过增强推理(TTA)机制对输入图像进行多次变换,并合并不同变换下的推理结果,提高了模型在测试时的鲁棒性和准确性。
  • 主要功能
    • 增强推理:通过尺度变化和翻转,生成多个预测结果,并将它们合并。
    • 还原预测:将变换后的预测结果还原到原图尺度,保证检测的正确性。
    • 结果合并:最终返回多次推理后的合并检测结果,提高模型的性能。

 4.4恢复图像尺寸

    def _descale_pred(self, p, flips, scale, img_size):
        """De-scales predictions from augmented inference, adjusting for flips and image size."""
        if self.inplace:
            p[..., :4] /= scale  # de-scale
            if flips == 2:
                p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
            elif flips == 3:
                p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
        else:
            x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
            if flips == 2:
                y = img_size[0] - y  # de-flip ud
            elif flips == 3:
                x = img_size[1] - x  # de-flip lr
            p = torch.cat((x, y, wh, p[..., 4:]), -1)
        return p

_descale_pred 函数解析

_descale_pred 函数的主要作用是在增强推理(Test Time Augmentation,TTA)中对预测结果进行反缩放反翻转,将经过尺度缩放和翻转后的预测结果恢复到原始输入图像的尺寸和方向。这是增强推理中至关重要的一步,因为在增强推理时图像经过缩放或翻转,预测的边界框也发生了相应的变化,需要通过这个函数还原到原始的图像尺度和方向。


功能概述

  • 反缩放:将经过缩放后的预测结果恢复到原图尺度。
  • 反翻转:如果图像在推理时进行了翻转,则需要反向翻转预测结果。
  • 参数解释
    • p:模型的预测结果(包含边界框的坐标、置信度、类别等)。
    • flips:翻转类型,用于指示推理时是否进行翻转。
      • 2 表示上下翻转。
      • 3 表示左右翻转。
    • scale:图像缩放因子(推理时使用的尺度因子,如 0.830.67)。
    • img_size:输入图像的尺寸(高度、宽度)。

代码逻辑解析

1. 判断是否使用 inplace 操作
if self.inplace:
  • inplace 是一个布尔变量,控制是否在原地修改张量。若为 True,则在原张量上进行修改;否则创建新的张量进行操作。

2. 反缩放处理
p[..., :4] /= scale  # de-scale
  • 功能:将预测结果中的前四个维度(通常是边界框的 x, y, w, h)进行反缩放,恢复到原始图像的尺寸。
    • p[..., :4]:表示边界框的四个参数:x(中心点横坐标)、y(中心点纵坐标)、w(宽度)、h(高度)。
    • p[..., :4] /= scale:对这些参数进行缩放还原。因为在增强推理时输入图像被缩放,所以推理出的边界框也被缩放过,现在通过除以缩放因子 scale 来还原到原图尺度。

3. 反翻转处理
if flips == 2:
    p[..., 1] = img_size[0] - p[..., 1]  # de-flip ud
elif flips == 3:
    p[..., 0] = img_size[1] - p[..., 0]  # de-flip lr
  • 功能:根据 flips 的值对预测结果进行翻转还原。
    • flips == 2:表示图像在推理时进行了上下翻转。通过计算 img_size[0] - p[..., 1],将边界框的 y 坐标重新翻转回来。
    • flips == 3:表示图像在推理时进行了左右翻转。通过 img_size[1] - p[..., 0],将 x 坐标重新翻转回来。

4. inplace 操作时的处理
x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale  # de-scale
  • 功能:如果不进行 inplace 操作,则将边界框的坐标(x, y)和尺寸(w, h)分别进行反缩放,创建新的张量来保存反缩放结果。
if flips == 2:
    y = img_size[0] - y  # de-flip ud
elif flips == 3:
    x = img_size[1] - x  # de-flip lr
  • 功能:同样的,进行翻转还原,计算新的 x, y 坐标,将上下或左右翻转操作还原。
p = torch.cat((x, y, wh, p[..., 4:]), -1)
  • 功能:将反缩放后的 x, y, w, h 与其他预测结果(如置信度和类别信息)拼接在一起,形成完整的预测结果。

小结

  1. 反缩放:根据推理时使用的缩放因子 scale,将预测的边界框恢复到原始图像的尺度。
  2. 反翻转:如果推理时对图像进行了翻转操作(上下或左右翻转),将预测结果的边界框进行相应的翻转还原。
  3. inplace 操作:如果启用了 inplace,则直接在原张量上进行操作,节省内存;否则创建新的张量进行操作。

4.5 TTA对原图片进行裁剪

    def _clip_augmented(self, y):
        """Clips augmented inference tails for YOLOv5 models, affecting first and last tensors based on grid points and
        layer counts.
        """
        nl = self.model[-1].nl  # number of detection layers (P3-P5)
        g = sum(4**x for x in range(nl))  # grid points
        e = 1  # exclude layer count
        i = (y[0].shape[1] // g) * sum(4**x for x in range(e))  # indices
        y[0] = y[0][:, :-i]  # large
        i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
        y[-1] = y[-1][:, i:]  # small
        return y

4.6初始化偏置biases信息

    def _initialize_biases(self, cf=None):
        """
        Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf).

        For details see https://arxiv.org/abs/1708.02002 section 3.3.
        """
        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
        m = self.model[-1]  # Detect() module
        for mi, s in zip(m.m, m.stride):  # from
            b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
            b.data[:, 5 : 5 + m.nc] += (
                math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
            )  # cls
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

_initialize_biases 函数解析

_initialize_biases 函数用于初始化 YOLOv5 模型中 Detect 层的偏置项(biases),以提高模型的预测能力,尤其是在训练初期。


函数的作用

该函数的主要目的是对 Detect 层中的卷积层的偏置进行初始化。具体来说,它对两种偏置进行初始化:

  1. 目标置信度(objectness score) 偏置。
  2. 类别分类(class scores) 偏置。

通过合理的偏置初始化,可以让模型在训练时更快收敛,尤其是在处理目标检测任务时,有效提升模型的初期性能。


关键参数

  • cf(可选参数):类别频率信息(class frequencies),通常用来表示每个类别在数据集中出现的频率。如果不传入 cf,则使用默认的初始化策略。
    • cf 的计算示例(注释中的代码):cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 这是一个计算数据集中每个类别的频率的方式。

核心流程

  1. 获取 Detect 层
m = self.model[-1]  # Detect() module
  • self.model[-1]:获取模型的最后一层 Detect 层,这是 YOLOv5 中专门用于预测边界框和类别的层。

  1. 遍历 Detect 层中的卷积模块和步幅(stride)
for mi, s in zip(m.m, m.stride):  # from
  • m.mDetect 层中的多个卷积层(通常每个特征图尺度上都有一个)。
  • m.stride:每个特征图对应的步幅(stride),表示特征图与原始图像的缩放倍数。步幅越大,特征图的分辨率越低,检测大物体的能力越强。

  1. 对偏置进行重塑和调整
b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
  • mi.bias.view(m.na, -1):将卷积层的偏置从扁平化的形式(如 255)重新塑形为 (3, 85),其中 3 表示 3 个锚框,85 表示每个锚框对应的输出数目(4 个坐标 + 1 个目标置信度 + 80 个类别概率,假设有 80 个类别)。
  • 通过 view 操作可以很方便地对偏置进行操作并更新。

  1. 初始化目标置信度的偏置
b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
  • b.data[:, 4]:指的是每个锚框输出的第 4 个位置,对应的是目标置信度(objectness score)的偏置。
  • 初始化策略:用 math.log(8 / (640 / s) ** 2) 来初始化目标置信度的偏置,这里的 8 表示每张 640×640 的图片上预期检测到的目标数目。这样可以让模型在训练的早期阶段对目标的置信度有较合理的初始化估计。

  1. 初始化类别分类的偏置
b.data[:, 5 : 5 + m.nc] += (
    math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
)  # cls
  • b.data[:, 5 : 5 + m.nc]:这是对类别分类概率的偏置进行初始化。
    • m.nc 表示类别的数量。
    • 偏置初始化策略
      • 如果 cf 为 None,使用默认的初始化:math.log(0.6 / (m.nc - 0.99999)),这会让模型在初始化时对每个类别有一个合理的起始概率。
      • 如果传入了 cf(类别频率),则通过类别频率 cf 来调整偏置项,使得频繁出现的类别的初始概率更高,稀有类别的初始概率更低。通过 torch.log(cf / cf.sum()) 计算。

  1. 将更新后的偏置重新赋值
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
  • 功能:将初始化后的偏置重新赋值回模型的卷积层中,并确保该偏置参与训练时的梯度更新(requires_grad=True)。
  • b.view(-1):将重新塑形的偏置恢复到扁平化形式,方便赋值给卷积层。

总结

  • 目标置信度偏置:根据图片大小和预期目标数进行初始化,使得模型在训练初期能够合理地估计目标置信度。
  • 类别分类偏置:根据类别数或类别频率进行初始化,帮助模型更快地学习类别信息,尤其是在类别分布不均衡的情况下。
  • 作用:该函数对偏置的初始化有助于加快模型的收敛速度,尤其是在目标检测任务中,可以显著提高模型在训练初期的性能。

4.5 ClassificationModel模块

class ClassificationModel(BaseModel):
    # YOLOv5 classification model
    def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):
        """Initializes YOLOv5 model with config file `cfg`, input channels `ch`, number of classes `nc`, and `cuttoff`
        index.
        """
        super().__init__()
        self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg)

    def _from_detection_model(self, model, nc=1000, cutoff=10):
        """Creates a classification model from a YOLOv5 detection model, slicing at `cutoff` and adding a classification
        layer.
        """
        if isinstance(model, DetectMultiBackend):
            model = model.model  # unwrap DetectMultiBackend
        model.model = model.model[:cutoff]  # backbone
        m = model.model[-1]  # last layer
        ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels  # ch into module
        c = Classify(ch, nc)  # Classify()
        c.i, c.f, c.type = m.i, m.f, "models.common.Classify"  # index, from, type
        model.model[-1] = c  # replace
        self.model = model.model
        self.stride = model.stride
        self.save = []
        self.nc = nc

    def _from_yaml(self, cfg):
        """Creates a YOLOv5 classification model from a specified *.yaml configuration file."""
        self.model = None

ClassificationModel 类解析

该类是 YOLOv5 中用于图像分类任务的模型结构。虽然 YOLOv5 最初是为目标检测设计的,但通过这个 ClassificationModel 类,可以将 YOLOv5 的特征提取能力用于图像分类任务。

主要功能
  • 将 YOLOv5 检测模型的前几层(主干网络部分)截取下来,用作分类任务的特征提取器。
  • 添加分类层,将模型输出转换为图像分类的类别预测。

1. 构造函数 __init__()

def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):
    """初始化 YOLOv5 分类模型
    - `cfg`: YOLOv5 配置文件(可选)
    - `model`: 已经构建好的 YOLOv5 模型(可选)
    - `nc`: 分类任务中的类别数,默认为 1000 类
    - `cutoff`: 截取检测模型的前 `cutoff` 层作为主干网络
    """
    super().__init__()
    self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg)
  • cfg:分类模型的配置文件。如果传入该参数,模型会根据 yaml 配置文件来创建。
  • model:已存在的 YOLOv5 检测模型。如果传入该参数,模型将基于现有的检测模型构建分类模型。
  • nc:分类任务中的类别数量,默认值为 1000(常用于 ImageNet 分类任务)。
  • cutoff:截断点,表示从检测模型中截取前 cutoff 层作为分类模型的特征提取部分。

2. _from_detection_model()

def _from_detection_model(self, model, nc=1000, cutoff=10):
    """从 YOLOv5 检测模型创建一个分类模型
    - `model`: YOLOv5 检测模型
    - `nc`: 类别数
    - `cutoff`: 截取点,用来保留检测模型的前 `cutoff` 层
    """
    if isinstance(model, DetectMultiBackend):
        model = model.model  # 如果传入的是 DetectMultiBackend,解包得到基础模型
    model.model = model.model[:cutoff]  # 截取主干网络(Backbone)
    m = model.model[-1]  # 获取最后一层
    ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels  # 获取最后一层的输入通道数
    c = Classify(ch, nc)  # 创建一个分类层,输入通道数为 `ch`,输出为 `nc` 类
    c.i, c.f, c.type = m.i, m.f, "models.common.Classify"  # 记录层的信息
    model.model[-1] = c  # 用分类层替换最后一层
    self.model = model.model  # 更新模型
    self.stride = model.stride  # 继承检测模型的 stride
    self.save = []  # 分类模型无需保存某些中间层的输出
    self.nc = nc  # 类别数

主要流程:

  • 检测模型解包:如果输入的模型是 DetectMultiBackend,则解包成普通的 YOLOv5 检测模型。
  • 截取主干网络:将检测模型的前 cutoff 层作为分类模型的主干部分,这部分相当于提取图像的特征(类似于 ResNet、VGG 的卷积层)。
  • 构建分类层
    • 获取主干网络最后一层的输入通道数 ch
    • 创建一个 Classify 层,输入通道为 ch,输出通道为 nc(即类别数)。
    • 将新创建的 Classify 层替换掉主干网络的最后一层(通常是用于目标检测的 Detect 层)。
  • 保存模型信息:更新模型,记录 stride 和类别数。

3. _from_yaml()

def _from_yaml(self, cfg):
    """根据指定的 `yaml` 配置文件创建一个 YOLOv5 分类模型"""
    self.model = None  # 尚未实现,从 yaml 文件创建分类模型的逻辑
  • 这个函数的作用是根据 yaml 配置文件创建一个分类模型。
  • 目前代码中尚未实现该功能。如果将 cfg 作为配置传入,则会调用此函数,但目前的实现并不会构建模型。

代码核心解析

  • ClassificationModel 类的核心目标是将 YOLOv5 的检测能力转换为图像分类能力。
  • 通过将原 YOLOv5 模型的检测部分裁剪掉,并添加一个新的分类层,可以快速将 YOLOv5 变为一个高效的图像分类器。
  • 分类层 Classify 是核心,它将特征提取后的信息转化为类别的概率分布。

总结

  1. 基于现有检测模型创建分类模型_from_detection_model() 函数通过裁剪检测模型的主干网络,并添加一个分类层来构建分类模型。
  2. 灵活的输入:可以选择从已有的检测模型或从配置文件中创建分类模型(尽管 yaml 的部分未实现)。
  3. 应用场景ClassificationModel 可用于图像分类任务,例如在 ImageNet 数据集上进行图像分类,也可以用于其它类似的分类任务。

五、总结

YOLOv5 yolo.py 文件的总结

在今天的讨论中,我们深入解析了 YOLOv5 的 yolo.py 文件,该文件是 YOLOv5 模型的核心实现之一,涵盖了模型的构建、推理、增强推理、分类任务的扩展等多个方面的功能。


1. 模型的基础构建

  • BaseModel 是 YOLOv5 模型的基础类,它定义了模型的通用功能,包括前向传播(推理或训练)、层性能分析、特征可视化、权重初始化和批量归一化层的融合等功能。

     

    关键方法:

    • forward(x):执行推理或训练操作。
    • _forward_once(x):执行单次推理的前向传播。
    • _profile_one_layer(m, x, dt):对网络的每一层进行性能分析,记录 FLOPs 和时间消耗。
    • fuse():将卷积层与批归一化层融合,提高推理速度。

2. 推理与增强推理(Test Time Augmentation, TTA)

  • _forward_augment():增强推理,用于测试时对图像进行不同尺度的缩放和翻转,然后合并多次推理结果以提高精度。

    • 增强推理的步骤
      1. 对输入图像进行尺度缩放和翻转。
      2. 使用 _forward_once 对每个变换后的图像进行推理。
      3. 使用 _descale_pred() 将推理结果还原到原图尺度,并合并所有结果。
  • _descale_pred():将增强推理后的预测结果反向缩放,并还原翻转操作,使预测结果能够与原始输入图像对齐。


3. 偏置的初始化

  • _initialize_biases():在 Detect 层中初始化 YOLOv5 的偏置项。通过合理的偏置初始化,模型可以在训练早期就能对目标置信度和类别有合理的估计,帮助模型更快收敛。
    • 置信度偏置:通过对目标置信度进行初始化,使得每张 640x640 的图像中预期检测到的目标数目为 8。
    • 类别偏置:可以根据类别频率(cf)来调整,使得频繁出现的类别有更高的初始概率。

4. 分类模型扩展

  • ClassificationModel 类:YOLOv5 的分类模型。尽管 YOLOv5 的主要功能是目标检测,通过将其检测模型的前几层(主干网络)截取下来,并添加一个分类层,可以将其转换为图像分类模型。
    • 核心步骤
      1. 截取 YOLOv5 检测模型的前 cutoff 层作为特征提取器。
      2. 添加分类层,输出类别的概率分布。
      3. 使用现有的 YOLOv5 模型主干进行图像分类任务。

5. 模型的加载与配置

  • parse_model():解析 YOLOv5 模型的配置(来自 YAML 文件),根据配置构建模型层,并初始化相关的参数。
  • DetectionModel 类:负责从 yaml 文件或字典配置中构建完整的 YOLOv5 检测模型。该类在模型初始化时,定义了各层的结构、步幅、锚框等信息。

6. YOLOv5 检测模型的关键模块

  • Detect 模块:YOLOv5 的检测头部模块,用于生成边界框、目标置信度和类别预测。该模块接受来自多个特征图层的输入,并将预测结果映射回原图的尺度。
    • 关键操作
      1. 生成特征图对应的网格(grid)并初始化锚框。
      2. 对预测结果进行解码,得到物体的位置和类别。
      3. 通过 _make_grid() 函数生成网格坐标,确保预测结果与特征图对齐。

总结

  • 模块化设计:YOLOv5 的 yolo.py 文件充分利用了模块化设计,模型的主干网络、分类头、检测头等部分可以根据需求灵活组合。
  • 增强推理的支持:通过 Test Time Augmentation (TTA),YOLOv5 能够在测试阶段通过多次推理来提高检测精度。这对提升模型性能尤为重要。
  • 分类任务扩展:除了目标检测,YOLOv5 还支持图像分类任务,用户可以通过裁剪检测模型的主干网络并添加分类层来实现。
  • 偏置初始化:偏置初始化策略有助于模型在训练初期对目标和类别的合理估计,提高训练的效率和准确性。

Top