YOLO v5 代码精读（3）YOLO网络结构(yolov5 test.py)

YOLO v5 代码精读（3）YOLO网络结构

推荐整理分享YOLO v5 代码精读（3）YOLO网络结构(yolov5 test.py)，希望有所帮助，仅作参考，欢迎阅读内容。

文章相关热门搜索词:yolov2代码详解,yolo v5精简代码,yolo v5精简代码,yolov2代码详解,yolov1代码,yolo v5精简代码,yolo v5精简代码,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# 默认anchoranchors:- [10,13, 16,30, 33,23] # P3/8- [30,61, 62,45, 59,119] # P4/16- [116,90, 156,198, 373,326] # P5/32

nc：表示检测的类别数量，这里默认取自coco数据集的80个类别

depth_multiple：控制网络的深度

width_multiple：控制网络的宽度

以上两个参数决定了yolo模型的规模。不同规模的yolo模型的yaml文件只有以上两个参数不同，其余参数完全一致。

anchors：这里的anchors默认为coco数据集中的anchors。anchors分为三栏，分别对应三个不同的feature map，也就是三个不同的检测层。三个检测层感受野不同，检测的物体大小也不同。感受野越大，检测的物体越大，anchor也就越大；反之同理。anchors会以二维数组的形式被加载。

网络结构# YOLOv5 v6.0 backbonebackbone:# [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 headhead:[[-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 模型代码解读

有关于各个模块的详细原理以及作用，请看我的另一篇博客：YOLOv5深度剖析

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: 输入图片的通道数 默认为3nc: 检测类别数量anchors: 表示anchor框"""# 父类的构造方法super().__init__()# 如果cf是加载好的字典结果if isinstance(cfg, dict):# 直接保存到模型中self.yaml = cfg # model dict# 若不是字典 则为yaml文件路径else: # is *.yamlimport yaml # for torch hub# 保存文件名self.yaml_file = Path(cfg).namewith open(cfg, errors='ignore') as f:# 将yaml文件加载为字典self.yaml = yaml.safe_load(f) # model dict# Define model# yaml.get('ch', ch)表示若不存在键'ch',则返回值chch = 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不为Noneif 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 stridem.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, Nonereturn self._forward_once(x, profile, visualize) # single-scale inference, traindef _forward_augment(self, x):# 这个技巧会将图片进行裁剪,并分别送入模型进行检测img_size = x.shape[-2:] # height, widths = [1, 0.83, 0.67] # scalesf = [None, 3, None] # flips (2-ud, 3-lr)y = [] # outputsfor 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]) # saveyi = self._descale_pred(yi, fi, si, img_size)y.append(yi)y = self._clip_augmented(y) # clip augmented tailsreturn torch.cat(y, 1), None # augmented inference, traindef _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 convolutiondef __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groupssuper().__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 convolutionsdef __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansionsuper().__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卷积，对数据进行降维，再进行常规卷积核的卷积。最后通过残差结构链接。Bottleneck既减少了参数量，又优化了计算，保持了原有的精度。

class Bottleneck(nn.Module):# Standard bottleneckdef __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansionsuper().__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 == c2def forward(self, x):return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))SPPF模块

空间金字塔池化。用在骨干网络收尾阶段，用于融合多尺度特征。

class SPPF(nn.Module):# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocherdef __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 dimensiondef __init__(self, dimension=1):super().__init__()self.d = dimensiondef 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 buildonnx_dynamic = False # ONNX export parameterdef __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layersuper().__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都要调用一次conv1x1self.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]) # convbs, _, 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] # xyy[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # whelse: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xywh = (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].deviceyv, 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