YOLO v5 代码精读（3）YOLO网络结构
YOLO模型共有五种模型规格，规格越大的模型准确率越高，相应的预测时间也就越长。一般默认选择YOLOv5s，也可根据需求选择更大或更小的模型。
这里以YOLO v5s为例，分析YOLO的网络结构。
YOLO v5模型思想
yolov5s.yaml文件
配置变量
# Parameters
# 检测的类别数
nc: 80 # number of classes
# 控制网络的深度 这两个参数决定了yolo模型的规模
depth_multiple: 0.33 # model depth multiple
# 控制网络的宽度
width_multiple: 0.50 # layer channel multiple
# 默认anchor
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32nc：表示检测的类别数量，这里默认取自coco数据集的80个类别
depth_multiple：控制网络的深度
width_multiple：控制网络的宽度
以上两个参数决定了yolo模型的规模。不同规模的yolo模型的yaml文件只有以上两个参数不同，其余参数完全一致。
anchors：这里的anchors默认为coco数据集中的anchors。anchors分为三栏，分别对应三个不同的feature map，也就是三个不同的检测层。三个检测层感受野不同，检测的物体大小也不同。感受野越大，检测的物体越大，anchor也就越大；反之同理。anchors会以二维数组的形式被加载。
网络结构
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
# [参数来自于哪一层, 模块中子模块重复的个数, 哪个模块, 模块创建时的参数]
# from: -1表示来自于上一层 [-1, 6]表示参数来自上一层与第6层
# number: 表示子模块重复的个数 这个数只有C3模块是不为1的，表示C3模块中，子模块重复三次
# module: 表示这一层是哪个模块
# args: 表示创建层的参数
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]模型结构分为两个部分：backbone与head，数组中的参数依次为[from, number, module, args]。backone与head中，每个层都是一个模块，而一个模块内又有多个层（这里的层指的是常规的卷积层或BN层等）。
from：表示参数来自与哪一层。-1表示来自于上一层，而数组形式表示参数来自多个层，类如[-1, 6]表示参数来自上一层与第6层。只有Concat层与最后的Detect的参数来自多个层，这是因为Contcat层要将多个层的参数叠在一起，Detect层要预测3个不同特征图的数据，这三个特征图分别来自17，20，23层。
number：表示模块内子模块重复的次数。只有C3是不为1的。C3模块由3个卷积层与一个Bottleneck层组成number为3时，就相当于1个C3模块中重复三次3个卷积层与一个Bottleneck层。
module：表示这一层使用哪个模块。
args：表示创建这一层模块的参数，args最终会被传参至模块的构造方法中。args再被传参前还会再做一次处理，用于控制网络的规模。
模型思想
文件中将yolo模型分为两个部分：1）backbone；2）head；
其中backbone指的是YOLO v5的骨干网络，用于提取图片中的特征；而head会从backbone的不同层次中获取特征，并作处理，最后分别从3个feature map（特征图）中得到预测结果。以下图示为YOLO v5的网络架构。为了方便表示，这里图中的模型为YOLO v5l，因为YOLOv5l中的depth_multiple与width_multiple均为1，方便表示。
观察YOLO模型的结构图，整个YOLO模型是一个全卷积网络。
backbone就是YOLOv5的骨干网络，骨干网络负责一直一直做卷积，一直向下提取特征，并将第4层，第6层，第8层的参数保留下来。
在Head中，完成了特征的向上融合以及向下融合，形成三个特征图。
从上到下观察第一个特征图，这个特征图的特征是由骨干网络的第四层特征，拼接上特征向上融合的特征，以及累积下来的细小特征。由于卷积层是用来提取轮廓特征的，卷积越多提取的特征越细致，通过特征向上融合，一些较大的轮廓特征占据了主要地位。且维度较高，参数最多，感受野比较大。因此第一个特征图预测大物体。
再看第二个特征图，它的特征来源于，第一个特征图再做卷积，进一步细化，再拼接上特征向上融合中期的特征，这个中期的特征源自于骨干网络第6层与第一次上采样的拼接，再做降维和卷积。因此第二个特征图的特征得到了进一步的细化，维度适中，参数适中，感受野适中。适合预测中型的物体。
最后是第三个特征图。第三个特征图的特征来自第二个特征图的进一步细化，并拼接上来自骨干网络最后一层再做卷积的特征。这一层的特征参数是最细致的，累积了前面大部分层的特征。且维度较低，参数也最受，感受野也比较低模式。适合预测小型物体。
YOLO 模型代码解读
Model类
这里只截取到到正向传播的部分
class Model(nn.Module):
"""
自定义YOLO模型类
继承torch.nn.Module类
"""
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
"""
cfg: YOLO v5模型配置文件 这里使用yolov5s模型
ch: 输入图片的通道数 默认为3
nc: 检测类别数量
anchors: 表示anchor框
"""
# 父类的构造方法
super().__init__()
# 如果cf是加载好的字典结果
if isinstance(cfg, dict):
# 直接保存到模型中
self.yaml = cfg # model dict
# 若不是字典 则为yaml文件路径
else: # is *.yaml
import yaml # for torch hub
# 保存文件名
self.yaml_file = Path(cfg).name
with open(cfg, errors='ignore') as f:
# 将yaml文件加载为字典
self.yaml = yaml.safe_load(f) # model dict
# Define model
# yaml.get('ch', ch)表示若不存在键'ch',则返回值ch
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
# 若类别数不为None且与yaml中的值不一致
if nc and nc != self.yaml['nc']:
# 给出提示
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
# 将yaml中的值修改为构造方法中的值
self.yaml['nc'] = nc # override yaml value
# 若anchors不为None
if anchors:
# 给出提示
LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
# 将yaml中的值改为构造方法中的值
self.yaml['anchors'] = round(anchors) # override yaml value
# 将模型的字典结构数据加载为真正的模型
# deepcopy()复杂产生一个新的对象
# model表示加载的模型 save表示需要保存参数的层的索引
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
# inplace指的是原地操作 如x+=1 有利于节约内存
self.inplace = self.yaml.get('inplace', True)
# Build strides, anchors
# 取出模型的最后一层 Detect层
m = self.model[-1] # Detect()
# 判断最后一层是否为Detect层
if isinstance(m, Detect):
s = 256 # 2x min stride
m.inplace = self.inplace
# 用一张大小为256的空图片进行一次前向传播
# 分别计算出 每张特征图上缩放了多少倍
# 计算出的stride为 [8, 16, 32]
# 第一个特征图上缩小了8倍，检测大物体 第二张特征图缩小了16倍，检测中物体 第三张特征图缩小了32倍，检测小物体
# 这里的缩小并不是指图片缩小了 而是指特征更细了
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
# 原始定义的anchor是原始图片上的像素值，要将其缩放至特征图的大小
m.anchors /= m.stride.view(-1, 1, 1)
# 检测anchor的顺序是否正确
check_anchor_order(m)
# 将步长保存至模型
self.stride = m.stride
# 初始化偏置参数
self._initialize_biases() # only run once
# Init weights, biases
# 初始化权重参数
initialize_weights(self)
# 设置日志信息
self.info()
LOGGER.info('')
def forward(self, x, augment=False, profile=False, visualize=False):
"""正向传播"""
# 是否使用数据增强
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):
# 这个技巧会将图片进行裁剪,并分别送入模型进行检测
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
def _forward_once(self, x, profile=False, visualize=False):
y, dt = [], [] # outputs
# 遍历模型中的每一层
for m in self.model:
# 判断数据的来源
if m.f != -1: # if not from previous layer
# 如果不是由上一层传递的参数,将从save的y中读取到,并改成列表的格式
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下面是模型中的几个自定义模块
Conv卷积模块
卷积完成后进行批归一化。
class Conv(nn.Module):
# Standard convolution
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super().__init__()
# c1为输入通道数 可以理解为卷积核的通道数
# c2为输出通道数 可以理解为卷积核的个数
# k为卷积核的大小
# s为卷积步长
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
# 批归一化层
self.bn = nn.BatchNorm2d(c2)
# SiLU为激活函数 act表示激活函数
self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
def forward(self, x):
# 参数经过卷积、批归一化最后由激活函数输出
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
# 不经过批归一化
return self.act(self.conv(x))C3模块
输入分成了两部分，一部分先经过cv1在经过self.m定义的Bottleneck操作，另一部分直接经过cv2，俩部分最后汇总拼接后一起经过cv3得到输出。该模块是对残差特征进行学习的主要模块,其结构分为两支,一支使用了上述指定多个Bottleneck堆叠和3个标准卷积层，另一支仅经过一个基本卷积模块，最后将两支进行concat操作。
class C3(nn.Module):
# CSP Bottleneck with 3 convolutions
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
# c1为输入通道数 可以理解为卷积核的通道数
# c2为输出通道数 可以理解为卷积核的个数
# c_为c2的一半
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
def forward(self, x):
# 将第一个卷积层与第二个卷积层的结果拼接在一起
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))Bottleneck类
Bottleneck结构就是为了降低参数量，Bottleneck先进行1x1卷积，对数据进行降维，再进行常规卷积核的卷积，是在Resnet中提出的。Bottleneck既减少了参数量，又优化了计算，保持了原有的精度。
class Bottleneck(nn.Module):
# Standard bottleneck
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
# 两个卷积模块
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c2, 3, 1, g=g)
# shortcut是将前层与后层特征相加
self.add = shortcut and c1 == c2
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))SPPF类
空间金字塔池化，能将任意大小的特征图转换成固定大小的特征向量。用在骨干网络收尾阶段，把特征变为固定大小送入head。
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
super().__init__()
c_ = c1 // 2 # hidden channels
# 两个卷积模块
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
# 最大池化层
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
# y1为最大池化的结果
y1 = self.m(x)
# 再进行一次最大池化
y2 = self.m(y1)
# 将所有结果拼到一起
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
Concat类
用于叠加多个层的参数
class Concat(nn.Module):
# Concatenate a list of tensors along dimension
def __init__(self, dimension=1):
super().__init__()
self.d = dimension
def forward(self, x):
# 将s上一层的参数叠加到一起
return torch.cat(x, self.d)nn.Upsample
上采样使用的是pytorch中实现的上采样上采样，又名放大图像、图像插值；
主要目的是放大原图像。上采样有3种常见的方法：双线性插值(bilinear)，反卷积(Transposed Convolution)，反池化(Unpooling)；图像放大几乎都是采用内插值方法，即在原有图像像素的基础上在像素点之间采用合适的插值算法插入新的元素。插值算法还包括了传统插值，基于边缘图像的插值，还有基于区域的图像插值。
Detect类
分别处理三个特征图的结果，并拼接到一起
class Detect(nn.Module):
# 特征图的缩放步长
stride = None # strides computed during build
onnx_dynamic = False # ONNX export parameter
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
super().__init__()
# 检测的类别数
self.nc = nc # number of classes
# 每个anchor输出的数据量 nc代表nc个类的条件类别概率 5包含4个定位参数和一个置信度参数
self.no = nc + 5 # number of outputs per anchor
# 特征图的数量
self.nl = len(anchors) # number of detection layers
# anchor的数量
self.na = len(anchors[0]) // 2 # number of anchors
# 初始化grid网格
self.grid = [torch.zeros(1)] * self.nl # init grid
# 初始化anchor网格
self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
# 将anchors注册到模型的buffers()属性中
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
# 对每个输出的feature map都要调用一次conv1x1
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
# 原地操作
self.inplace = inplace # use in-place ops (e.g. slice assignment)
def forward(self, x):
z = [] # inference output
# 对三个feature map分别进行处理
for i in range(self.nl):
# 进行1*1卷积
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
# 构造网格
# 因为推理返回的不是归一化后的网格偏移量 需要再加上网格的位置 得到最终的推理坐标 再送入nms
# 所以这里构建网格就是为了纪律每个grid的网格坐标 方面后面使用
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
# sigmoid激活函数
y = x[i].sigmoid()
# 计算定位参数
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
# cat(dim=-1)为直接拼接
y = torch.cat((xy, wh, y[..., 4:]), -1)
# 保存每个特征图的结果
z.append(y.view(bs, -1, self.no))
# 训练阶段直接返回x
# 预测阶段返回3个特征图拼接的结果
return x if self.training else (torch.cat(z, 1), x)
def _make_grid(self, nx=20, ny=20, i=0):
d = self.anchors[i].device
yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])
grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
.view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
return grid, anchor_grid