前言
YOLOv5就像一座金矿，里面有无数可以学习的东西。之前的博文一直将YOLOv5当作一个黑盒使用，只考虑模型的输入和输出，以此来对模型进行二次开发。
本篇博文将更近一层，深入到“金矿”内部，来尝试对模型结构进行替换。
模型构建解析
YOLOv5是通过yaml格式的模型配置文件来搭建模型架构的，这里我之前的博文【目标检测】YOLOv5：模型构建解析已经做过了解读，对此不再复述。
YOLOv5模型主要分5.0和6.0及以上版本，两者有少许区别，本文以后者模型为主。
YOLOv5s模型架构图如下，此图来源于目标检测 YOLOv5网络v6 0版本总结
修改模型
本文修改的目标是修改18、21这两个卷积块，这里是通过一个卷积核为3，步长为2的卷积核实现下采样，我的目标是修改为两个不同尺寸的卷积核，输出结果为两个不同卷积核之和。
验证维度
修改尺寸最麻烦的就是维度变化，因此在修改之前，最好对修改的部分单独模拟数据查看shape。
下面是一个测试示例：
import torch.nn as nn
import torch
def autopad(k, p=None, d=1): # kernel, padding, dilation
# Pad to 'same' shape outputs
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
class Conv(nn.Module):
# Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = self.default_act 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))
class Multi_Conv(nn.Module):
# Multi Different Kernel-size Conv
def __init__(self, c1, c2, e=1.0):
super().__init__()
c_ = int(c2 * e)
self.cv1 = Conv(c1, c_, 3, 2)
self.cv2 = Conv(c1, c_, 7, 2)
def forward(self, x):
return self.cv1(x) + self.cv2(x)
if __name__ == '__main__':
input_tensor = torch.rand(1, 128, 80, 80)
conv = Conv(128, 256, 3, 2)
mult_conv = Multi_Conv(128, 256)
output_tensor1 = conv(input_tensor)
print(output_tensor1.shape) # torch.Size([1, 256, 40, 40])
output_tensor2 = mult_conv(input_tensor)
print(output_tensor2.shape) # torch.Size([1, 256, 40, 40])
注：Conv并非pytorch原生的卷积，yolov5作者对其进行了重构，添加了autopad这个函数，这个可以让人在修改卷积核大小时，自动填充padding，以保证输出结果维度一致。
从上面这个示例可知，添加了我原创的双卷积核结构Multi_Conv之后，输出维度和单核输出一致。
嵌入模型
修改模型主要有两个方法，第一种是直接修改配置文件(.yaml)，yaml主要是用来控制模型的串行连接，修改完之后意味着后面的标号也需要进行调整，较为麻烦。
另一种思路就是模块替代，在模型单核模块中，替换成一个复杂的结构，这里选择第二种方法。
首先将创建的原创结构添加到models/common.py文件中：
class Multi_Conv(nn.Module):
# Multi Different Kernel-size Conv
def __init__(self, c1, c2, e=1.0):
super().__init__()
c_ = int(c2 * e)
self.cv1 = Conv(c1, c_, 3, 2)
self.cv2 = Conv(c1, c_, 7, 2)
def forward(self, x):
return self.cv1(x) + self.cv2(x)
然后在models/yolo.py中添加Multi_Conv：
if m in {
Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x, Multi_Conv}
添加完成之后，运行一下yolo.py，可以看到自己创立的模块已经被成功加载：
查看速度和参数量
在设计网络模型时，最好能够直观查看模型各层运行效率，在yolo.py中，作者预留了line-profile这个参数接口，设为True之后，可以看到模型每一层的参数量用时：
time (ms) GFLOPs params module
6.75 0.73 3520 models.common.Conv
0.70 0.96 18560 models.common.Conv
2.09 0.98 18816 models.common.C3
0.54 0.95 73984 models.common.Conv
1.86 1.49 115712 models.common.C3
0.40 0.95 295424 models.common.Conv
2.59 2.01 625152 models.common.C3
0.60 0.95 1180672 models.common.Conv
1.40 0.95 1182720 models.common.C3
0.60 0.53 656896 models.common.SPPF
0.20 0.11 131584 models.common.Conv
0.10 0.00 0 torch.nn.modules.upsampling.Upsample
0.00 0.00 0 models.common.Concat
1.50 1.16 361984 models.common.C3
0.30 0.11 33024 models.common.Conv
0.00 0.00 0 torch.nn.modules.upsampling.Upsample
0.00 0.00 0 models.common.Concat
1.40 1.17 90880 models.common.C3
6.65 3.04 950784 models.common.Multi_Conv
0.00 0.00 0 models.common.Concat
1.40 0.95 296448 models.common.C3
4.19 1.52 1901056 models.common.Multi_Conv
0.00 0.00 0 models.common.Concat
1.30 0.90 1117184 models.common.C3
0.50 0.73 229245 Detect
35.05 - - Total