一、导入包和本地配置
import argparse import contextlib import math import os import platform import sys from copy import deepcopy from pathlib import Path import torch import torch.nn as nn
argparse:
它是一个用于命令行选项和参数解析的模块。通过在程序中定义好我们需要的参数,argparse
将会从sys.argv
中解析出这些参数,并自动生成帮助和使用信息。contextlib:
它提供了用于简化上下文管理器使用的工具。上下文管理器主要用于管理诸如文件打开/关闭这类资源的分配和释放,contextlib
模块可以简化这些操作。math:
它是Python的数学库,提供了常用的数学函数和常量,如平方根、对数、三角函数等。os:
该模块提供了一些与操作系统交互的函数,比如文件和目录的操作、系统环境变量的获取等。它在跨平台编程中非常有用,可以使代码与操作系统无关。platform:
该模块提供了获取操作系统和Python版本信息的功能。通过它可以检测当前代码的运行环境,便于编写跨平台代码。sys:
它是Python解释器本身的一个接口,提供了与解释器交互的方法和函数。sys
模块可以访问和修改Python运行时环境的变量和函数。copy.deepcopy:
它是Python中用于复制对象的库。deepcopy
可以复制一个对象及其包含的所有对象,即使是嵌套结构。这样复制出的对象与原始对象完全独立。pathlib:
这个库提供了一种面向对象的方式来操作文件和目录路径。与传统的os.path
不同,pathlib
提供了更直观、简洁的文件路径操作方式,使代码更易读和维护。
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的网络结构,包括C3
、SPPF
、GhostConv
等模块,这些模块用于构建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
中提取模型的相关参数,并为后续的网络构建做好准备。具体步骤如下:
打印模型结构表头:
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
- 这行代码通过
LOGGER
打印模型结构表头,包括 "from"、"n"、"params"、"module" 和 "arguments" 这些列,用于后续日志输出中显示模型的层次结构。提取配置参数:
anchors, nc, gd, gw, act, ch_mul = ( d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"], d.get("activation"), d.get("channel_multiple"), )
- 从配置字典
d
中提取模型相关的关键参数:
anchors
:锚框的配置。nc
:类别数量(number of classes)。gd
:深度倍数(depth_multiple)。gw
:宽度倍数(width_multiple)。act
:激活函数(如果有的话)。ch_mul
:通道倍数(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
- 如果配置中指定了激活函数,则通过
eval
函数将其转为相应的激活函数,并将其设置为Conv
模块的默认激活函数,同时在日志中打印出当前使用的激活函数。确定通道倍数的默认值:
if not ch_mul: ch_mul = 8
- 如果配置中没有提供
channel_multiple
参数,则设置一个默认值8
。计算锚框数量:
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
- 计算锚框的数量。如果
anchors
是一个列表,则通过计算每个锚框的宽度和高度(即除以2)来确定数量。计算输出通道数:
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
- 计算每层的输出通道数,即
no
。输出通道数等于锚框数量乘以(类别数nc
加上5个参数,分别是4个边界框参数和1个置信度)。
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
这段代码主要是遍历
backbone
和head
的每一层,获取搭建网络结构前的一系列信息。我们还是先解释参数
layers
,save
和c2
:
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
:
当前层的模块类型,如Conv
、Focus
等,这些模块会被动态加载。
args
:
表示当前层的参数列表,例如卷积层的输出通道数、卷积核大小等。
eval()
函数的使用接着遇到一个函数
eval()
,主要用于将字符串转换成有效的表达式来求值,并且返回执行的结果。m = eval(m) if isinstance(m, str) else m # eval strings
m = eval(m)
:
- 如果
m
是字符串(即层的模块类型以字符串形式给出),通过eval
将字符串转为实际的模块对象(如Conv
,C3
等)。
同样地,
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
(当前层的输出通道数):
n = max(round(n * gd), 1)
:
n
是当前层的深度系数,通过乘以gd
(depth_multiple)来调整。gd
控制模型的深度,例如 YOLOv5s 的gd
值为 0.33,意味着把模型的深度缩减为原来的 1/3。round()
用来取整,并且至少为 1 层(避免n=0
的情况)。判断模块是否属于指定类型:
if m in {Conv, GhostConv, Bottleneck, GhostBottleneck, ... }:
- 判断当前层
m
是否属于卷积类模块,如Conv
、C3
、SPP
、Focus
等,这些模块通常会涉及输入通道和输出通道的变换。提取当前层的输入通道数
c1
和输出通道数c2
:c1, c2 = ch[f], args[0]
c1
:当前层的输入通道数,ch[f]
表示从前面第f
层输入。c2
:当前层的输出通道数,args[0]
是该层的输出通道数。调整输出通道数
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)
确保输出通道数c2
是ch_mul
的倍数,使其适合硬件加速和模型优化。下面是make_divisible
的实现:def make_divisible(x, divisor): return math.ceil(x / divisor) * divisor
make_divisible
:
- 作用是让
x
能够被divisor
整除,通过向上取整,使x
成为divisor
的最小倍数。- 这种方式确保计算效率,并在某些硬件平台上优化性能。
为了更好地理解这段代码的实际运行过程,以下是一个简化的运行案例,展示如何动态调整YOLOv5模型的深度和宽度,并解释各个部分是如何工作的。
案例背景
假设我们正在构建YOLOv5的一个卷积层模块,例如
Conv
或C3
,并且我们希望根据模型配置中的深度和宽度系数动态调整模型的深度和输出通道数。输入数据
我们有一个
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_multiple
和width_multiple
控制模型的深度和宽度。模拟模型层
假设我们正在处理一个
C3
模块,模块参数如下:n = 3 # 层的重复次数,通常是一个初始值,表示该层中子层的个数 ch = [64, 128, 256] # 通道数列表,表示前几层的输出通道数 args = [256] # 目标输出通道数 f = 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 次,模型变得更浅,减少计算量。
- 检查是否属于调整类型的层:
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
。
- 输出通道调整(
c2
的变化):gw = 0.5 # width_multiple 控制模型宽度 c2 = make_divisible(c2 * gw, 8) # 调整c2
- 计算
c2
:c2 = 256 * 0.5 = 128
。- 然后通过
make_divisible
确保c2
是 8 的倍数:c2 = make_divisible(128, 8) # 结果仍然是 128
- 这里没有变化,因为128已经是8的倍数。
- 结果输出:
最终,我们调整后的层结构如下:
- 输入通道数:
c1 = 128
- 输出通道数:
c2 = 128
- 深度:层重复 1 次,而不是 3 次。
这意味着,原本计划输出 256 个通道并重复 3 次的层,现在输出 128 个通道,并且只重复 1 次。通过这种方式,模型被动态缩小,适应资源受限的环境(例如移动设备)。
完整流程总结
- 深度调整:通过
depth_multiple
(gd
),减少某些层的深度,降低计算开销。- 宽度调整:通过
width_multiple
(gw
),减少每层的输出通道数,使模型变窄,减少参数量和计算量。- 输出调整:通过
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]
这段代码的主要功能是根据当前层的参数构建当前层的网络结构,并且为不同类型的模块(如
BottleneckCSP
、Concat
、Detect
等)动态计算输入和输出通道数。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
参数列表的前面。
- 如果当前层属于
BottleneckCSP
、C3
等模块,表示这类模块具有重复结构,需要将重复次数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
。特定模块的逻辑:
- 对于卷积类模块(如
BottleneckCSP
、C3
),输出通道数由配置中的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
是模块的类型名称,如Conv
,C3
,SPPF
等。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
):将处理好的xy
、wh
和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介绍
这几个模型类(
BaseModel
、DetectionModel
、SegmentationModel
、ClassificationModel
)是 YOLOv5 框架中的核心组件,分别用于处理不同的任务:目标检测、语义分割和图像分类。每个模型类有其特定的用途,能够解决不同的计算机视觉任务。下面是每个模型的具体用途和功能:
1. BaseModel
用途:这是 YOLOv5 中所有模型的基类,定义了一些通用的功能,比如前向传播、层的性能分析、模型信息打印、层融合等。
BaseModel
为其他模型(如DetectionModel
、SegmentationModel
、ClassificationModel
)提供了核心的基础功能。
- 典型用途:
- 为 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(浮点运算次数)、执行时间和参数数量。对于每一层都会记录这些信息,便于优化和调试。
- 流程:
- 使用
thop
库计算 FLOPs(每秒浮点运算次数)。- 使用
time_sync()
记录每层的执行时间,循环运行 10 次,取平均值。- 打印每一层的参数数量、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
层,会进一步对stride
、grid
和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 模型在推理和训练时能够正确计算边界框的位置和大小。接下来,我将逐步解析这些代码的实现及其意义。
主要步骤:
获取模型的最后一层
Detect
层m = self.model[-1] # Detect()
self.model[-1]
:表示获取模型的最后一层(一般是Detect
层),这层用于处理目标检测的最终输出。Detect
层:YOLOv5 用于预测目标边界框、置信度和类别的核心层。
定义一个 256×256 大小的输入
s = 256 # 2x min stride
- 定义
s = 256
:表示模型的输入大小为 256×256,这个大小用于模拟推理时的输入图像尺寸。- 为什么使用 256:这是为了计算每一层输出特征图相对于输入图像的缩放比例,即步幅(stride)。通过定义一个固定大小的输入,模型可以推理出每个检测层的特征图步幅。
通过一次前向传播计算步幅
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
:步幅张量,存储了每层特征图的缩放倍数(步幅)。
调整锚框大小
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
中,供后续使用。
初始化偏置
self._initialize_biases() # only run once
- 功能:初始化
Detect
层的偏置项(bias),这个步骤仅会在模型构建时运行一次。- 作用:通过初始化偏置项,使模型在训练时能够更快速地收敛,特别是针对目标检测任务的边界框预测。
代码的整体流程总结
- 获取最后一层的
Detect
层:用于目标检测的输出层。- 定义 256×256 大小的输入:作为示例输入,模拟真实图像的推理过程。
- 通过一次前向传播计算特征图的步幅:利用输入图像大小和输出特征图大小的比例,计算每层特征图的缩放倍数(stride)。
- 调整锚框大小:根据每层特征图的步幅调整锚框的大小,使得锚框能适应不同尺度的特征图。
- 初始化偏置:确保偏置的合理初始化,提高模型的训练效果。
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, 0.83, 0.67]
,分别表示原始尺寸、83% 缩小和67% 缩小,同时设置了翻转选项[None, 3, None]
,其中3
代表左右翻转。- 循环处理:遍历每个尺度和翻转组合,先对输入图像进行缩放和翻转,然后调用
_forward_once
进行推理。- 反缩放与还原预测:通过
_descale_pred
函数将预测结果还原回原图尺度。- 结果合并:将不同尺度的推理结果存储在列表
y
中,最后使用torch.cat
将它们沿着通道维度进行拼接。- 裁剪:通过
_clip_augmented
函数对增强推理的结果进行裁剪,去除不必要的尾部。
增强推理的详细步骤
增强推理(TTA,Test Time Augmentation)是为了在推理阶段通过多种数据变换提高模型的鲁棒性和检测准确率。具体来说:
输入图像变换:
- 图像在原始尺度的基础上进行不同的缩放:
1
表示原始大小,0.83
表示缩小至原尺寸的83%,0.67
表示缩小至原尺寸的67%。- 同时使用翻转:
None
表示不翻转,3
表示水平翻转(左右翻转)。多次推理:
- 对每次变换后的图像进行单次推理(通过
_forward_once
函数)。- 推理的结果为检测到的边界框、置信度和类别标签。
还原预测结果:
- 每次推理的结果都需要进行反缩放和翻转还原,使得这些结果能够和原始图像对齐。
- 通过
_descale_pred
函数对预测结果进行还原。合并结果:
- 将不同变换下的推理结果拼接在一起,通过
torch.cat(y, 1)
将所有推理结果在通道维度上进行合并。裁剪:
- 为了去除多余的边界框预测结果,使用
_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.83
、0.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
与其他预测结果(如置信度和类别信息)拼接在一起,形成完整的预测结果。
小结
- 反缩放:根据推理时使用的缩放因子
scale
,将预测的边界框恢复到原始图像的尺度。- 反翻转:如果推理时对图像进行了翻转操作(上下或左右翻转),将预测结果的边界框进行相应的翻转还原。
- 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
层中的卷积层的偏置进行初始化。具体来说,它对两种偏置进行初始化:
- 目标置信度(objectness score) 偏置。
- 类别分类(class scores) 偏置。
通过合理的偏置初始化,可以让模型在训练时更快收敛,尤其是在处理目标检测任务时,有效提升模型的初期性能。
关键参数
cf
(可选参数):类别频率信息(class frequencies),通常用来表示每个类别在数据集中出现的频率。如果不传入cf
,则使用默认的初始化策略。
cf
的计算示例(注释中的代码):cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
这是一个计算数据集中每个类别的频率的方式。
核心流程
- 获取
Detect
层m = self.model[-1] # Detect() module
self.model[-1]
:获取模型的最后一层Detect
层,这是 YOLOv5 中专门用于预测边界框和类别的层。
- 遍历
Detect
层中的卷积模块和步幅(stride)for mi, s in zip(m.m, m.stride): # from
m.m
:Detect
层中的多个卷积层(通常每个特征图尺度上都有一个)。m.stride
:每个特征图对应的步幅(stride),表示特征图与原始图像的缩放倍数。步幅越大,特征图的分辨率越低,检测大物体的能力越强。
- 对偏置进行重塑和调整
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
操作可以很方便地对偏置进行操作并更新。
- 初始化目标置信度的偏置
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 的图片上预期检测到的目标数目。这样可以让模型在训练的早期阶段对目标的置信度有较合理的初始化估计。
- 初始化类别分类的偏置
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())
计算。
- 将更新后的偏置重新赋值
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
是核心,它将特征提取后的信息转化为类别的概率分布。
总结
- 基于现有检测模型创建分类模型:
_from_detection_model()
函数通过裁剪检测模型的主干网络,并添加一个分类层来构建分类模型。- 灵活的输入:可以选择从已有的检测模型或从配置文件中创建分类模型(尽管
yaml
的部分未实现)。- 应用场景:
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()
:增强推理,用于测试时对图像进行不同尺度的缩放和翻转,然后合并多次推理结果以提高精度。
- 增强推理的步骤:
- 对输入图像进行尺度缩放和翻转。
- 使用
_forward_once
对每个变换后的图像进行推理。- 使用
_descale_pred()
将推理结果还原到原图尺度,并合并所有结果。
_descale_pred()
:将增强推理后的预测结果反向缩放,并还原翻转操作,使预测结果能够与原始输入图像对齐。
3. 偏置的初始化
_initialize_biases()
:在Detect
层中初始化 YOLOv5 的偏置项。通过合理的偏置初始化,模型可以在训练早期就能对目标置信度和类别有合理的估计,帮助模型更快收敛。
- 置信度偏置:通过对目标置信度进行初始化,使得每张 640x640 的图像中预期检测到的目标数目为 8。
- 类别偏置:可以根据类别频率(
cf
)来调整,使得频繁出现的类别有更高的初始概率。
4. 分类模型扩展
ClassificationModel
类:YOLOv5 的分类模型。尽管 YOLOv5 的主要功能是目标检测,通过将其检测模型的前几层(主干网络)截取下来,并添加一个分类层,可以将其转换为图像分类模型。
- 核心步骤:
- 截取 YOLOv5 检测模型的前
cutoff
层作为特征提取器。- 添加分类层,输出类别的概率分布。
- 使用现有的 YOLOv5 模型主干进行图像分类任务。
5. 模型的加载与配置
parse_model()
:解析 YOLOv5 模型的配置(来自 YAML 文件),根据配置构建模型层,并初始化相关的参数。DetectionModel
类:负责从yaml
文件或字典配置中构建完整的 YOLOv5 检测模型。该类在模型初始化时,定义了各层的结构、步幅、锚框等信息。
6. YOLOv5 检测模型的关键模块
Detect
模块:YOLOv5 的检测头部模块,用于生成边界框、目标置信度和类别预测。该模块接受来自多个特征图层的输入,并将预测结果映射回原图的尺度。
- 关键操作:
- 生成特征图对应的网格(grid)并初始化锚框。
- 对预测结果进行解码,得到物体的位置和类别。
- 通过
_make_grid()
函数生成网格坐标,确保预测结果与特征图对齐。
总结
- 模块化设计:YOLOv5 的
yolo.py
文件充分利用了模块化设计,模型的主干网络、分类头、检测头等部分可以根据需求灵活组合。- 增强推理的支持:通过
Test Time Augmentation (TTA)
,YOLOv5 能够在测试阶段通过多次推理来提高检测精度。这对提升模型性能尤为重要。- 分类任务扩展:除了目标检测,YOLOv5 还支持图像分类任务,用户可以通过裁剪检测模型的主干网络并添加分类层来实现。
- 偏置初始化:偏置初始化策略有助于模型在训练初期对目标和类别的合理估计,提高训练的效率和准确性。