YOLOv5的模型构建源码详解|CSDN创作打卡
深度学习入门小菜鸟，希望像做笔记记录自己学的东西，也希望能帮助到同样入门的人，更希望大佬们帮忙纠错啦~侵权立删。
代码分析注释全家桶部分只是为了方便看循环，条件判断的那些缩进对应，与二、三、四讲的东西是一样的。可以直接粘贴进VS code里面看，也会比较清晰，看你喜欢什么样的形式啦~
目录
一、碎碎念前情提要
二、def __init__
1、加载解析yaml配置文件（包含网络参数等等）
2、self.yaml
3、定义网络模型
4、Build strides, anchors
5、初始化权重和偏置
三、def forward
1、主体
2、_forward_once函数
五、代码分析注释全家桶
1、Model部分
2、_forward_once函数
3、parse_model函数
一、碎碎念前情提要
在前面我们已经介绍了YOLOv5的网络框架等，传送门
YOLOv5网络结构详解_tt丫的博客-CSDN博客
那要怎么把它具体地搭建起来呢？那就让我们从源码（YOLOv5源码中的yolo.py的class Model）入手吧
二、def __init__
其实跑完__init__，整个网络就搭起来了。
1、加载解析yaml配置文件（包含网络参数等等）
class Model(nn.Module):
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
super().__init__()
if isinstance(cfg, dict): #判断cfg是不是dict（字典）类型
self.yaml = cfg # model dict
#模型词典赋值给self.yaml首先判断cfg是不是dict（字典）类型，如果是，就把这个模型词典赋值给self.yaml
else: # is *.yaml
import yaml # for torch hub
self.yaml_file = Path(cfg).name # 获取cfg的文件名
with open(cfg, encoding='ascii', errors='ignore') as f:#用ascii编码，忽略错误的形式打开文件cfg
self.yaml = yaml.safe_load(f) # model dict
#用yaml的文件形式加载cfg，赋值给self.yaml否则导入yaml，获取cfg的文件名后，用ascii编码，忽略错误的形式打开文件cfg，用yaml的文件形式加载cfg，赋值给self.yaml。
最后的yaml就是字典形式的配置信息。
2、self.yaml
让我们先来看看配置文件里的都是些啥
首先是一些参数的设定
# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32然后是网络backbone部分
# YOLOv5 v6.0 backbone
backbone:
# [from, number, 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
]最后是网络Neck + Head部分（不过这里它是直接写成Head了，其实都一样）
# 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)
]3、定义网络模型
# 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判断输入的channel和配置文件里的是否相同，不相同则变成输入的参数。
if anchors: LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}') self.yaml['anchors'] = round(anchors) # override yaml value将anchors进行四舍五入（防止输入的是小数从而报错）
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist将配置文件中的信息导入模型中并保存。
然后让我们穿插一下这里用到的parse_model函数
请看目录
四、其他所需函数 —— parse_model函数
让我们重新回去继续分析
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
#给那些种类编编号，从0到nc-1
self.inplace = self.yaml.get('inplace', True)self.name：给那些种类编编号，从0到nc-1
4、Build strides, anchors
m = self.model[-1] # Detect()
if isinstance(m, Detect):
s = 256 # 2x min stride
m.inplace = self.inplace
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
m.anchors /= m.stride.view(-1, 1, 1)
这里调用了forward()函数：
然后让我们穿插一下这里用到的forward（）函数
请看目录
三、def forward
让我们回来继续分析。
输入一个[1, ch, 256, 256]的tensor，接着获得FPN输出结果的维度。然后求出了下采样的倍数stride：8，16，32。
最后用anchor去除以上的数值，这样就把anchor放缩到了3个不同的尺度上。
因此anchor的最终shape是[3,3,2]。
check_anchor_order(m)
#根据YOLOv5 Detect（）模块m的步幅顺序检查锚定顺序，必要时进行纠正
self.stride = m.stride
self._initialize_biases() # only run once
#将偏差初始化进Detect模块首先根据YOLOv5 Detect（）模块m的步幅顺序检查锚定顺序，必要时进行纠正，检查纠正后stride就定下来了，然后将偏差初始化进Detect模块（没有偏差，因为初始类频率cf为None）。
5、初始化权重和偏置
initialize_weights(self)
self.info()
LOGGER.info('')三、def forward
1、主体
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如果augment是True就调用_forward_augment函数，做增强处理。
如果augment是False就调用_forward_once函数，不做增强处理。
可以看到默认augment为False，所以其实实际上调用的是_forward_once函数。
2、_forward_once函数
def _forward_once(self, x, profile=False, visualize=False):
y, dt = [], [] # outputs
for m in self.model:这里先初始化存储输出的列表（其中dt没有用到，后续会说明），然后进入循环，遍历每个模块（层）
再次说明一下m中的参数：m.f是从哪层开始；
m.n是模块的默认深度；
m.args是该层的参数（就是从yaml那来的）
if m.f != -1: # if not from previous layer
#不是上一层,就是说不直接连接上一层（比如像Concat)
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
#取出对应的层的结果，准备后面的进入对应m的forward()首先先判断一下：如果不直接连接上一层，就比如说Concat，那么我们取出对应的层的结果，准备后面的进入对应m的forward()。
（对应的层的结果和对应m的forward()后面会说）
if profile:
self._profile_one_layer(m, x, dt)因为默认的profile是False，输入的也是False，所以没有用到，咱不说他，下一个~
x = m(x) # run
#m是模块（某个层）的意思，所以x传入模块，相当于执行模块（比如说Focus，SPP等）中的forward（）
#第一层Focus的m.f是-1，所以直接跳到这一步开始这里就是我们前面所说的“对应m的forward()”啦
m是模块（某个层）的意思，所以x传入模块，相当于执行模块（比如说Focus，SPP等）中的forward（）
因为第一层Focus的m.f是-1，前面的那些判断一个没进，所以直接跳到这一步开始。
y.append(x if m.i in self.save else None) # save output
#将每一层的输出结果保存到y这里就是我们前面所说的“对应的层的结果”啦
这里是将每一层的输出结果保存到y，这里的保存是给像Concat这些模块服务的，他们就不是直接连接上一层的，需要全部待连接层的输出结果。
因果关系嘛，所以不是所有层的输出结果都要保存，只有被需要，才要保存。
这里就跟我们的self.m有关系啦。如果你还记得的话，其实我们前面说过的，在parse_model函数中的这里
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist这里就有按需分配的操作啦（x!=1）
好的回归我们的_forward_once函数
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)这里跟profile一样，默认False，不需要用到，pass~
return xreturn就完了
四、其他所需函数 —— parse_model函数
def parse_model(d, ch): # model_dict, input_channels(3)
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
#日志记载，不管他
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)以上是读取配置dict里的参数，具体参数都是些啥可以看我之前的文章
YOLOv5的输出端（Head）详解|CSDN创作打卡_tt丫的博客-CSDN博客
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 # 将模块类型m转为值
for j, a in enumerate(args):#循环模块参数args
try:
args[j] = eval(a) if isinstance(a, str) else a # 将每个模块参数args[j]转为值
except NameError:
pass
迭代循环backbone与Head（Neck，head）的配置，并且将模块类型m转为值，每个模块参数args[j]转为值
f：从哪层开始
n：模块的默认深度
m：模块的类型
args：模块的参数
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
#网络用（n*gd）控制模块的深度缩放。
#深度在这里指的是像CSP这种模块的重复迭代次数，宽度一般我们指的是特征图的channel。网络用（n*gd）控制模块的深度缩放。
深度在这里指的是像CSP这种模块的重复迭代次数，宽度一般我们指的是特征图的channel。
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
c1, c2 = ch[f], args[0]
#ch是用来保存之前所有的模块输出的channle（所以ch[-1]代表着上一个模块的输出通道）。
# ch[f]是第f层的输出。args[0]是默认的输出通道。ch是用来保存之前所有的模块输出的channle（所以ch[-1]代表着上一个模块的输出通道）。
ch[f]是第f层的输出。args[0]是默认的输出通道。
if c2 != no: # if not output
#如果不是最终的输出
c2 = make_divisible(c2 * gw, 8)#保证了输出的通道是8的倍数
def make_divisible(x, divisor):
# Returns nearest x divisible by divisor
#返回可被除数整除的最近x
if isinstance(divisor, torch.Tensor):
divisor = int(divisor.max()) # to int
return math.ceil(x / divisor) * divisor如果不是最终的输出，调用general.py中的make_divisible函数，生成可被8整除的最近x作为c2，这样就保证了输出的通道是8的倍数。
args = [c1, c2, *args[1:]] #args变为原来的args+module的输入通道数（c1）、输出通道数（c2）
if m in [BottleneckCSP, C3, C3TR, C3Ghost]: #只有CSP结构的才会根据深度参数n来调整该模块的重复迭加次数
args.insert(2, n) # number of repeats
#模块参数信息args插入n
n = 1#n重置对args进行信息补充：
args变为原来的args+module的输入通道数（c1）、输出通道数（c2）
如果是BottleneckCSP, C3的话，还要加上深度参数n（用来调整该模块的重复迭加次数），然后重置n。
elif m is nn.BatchNorm2d:
args = [ch[f]]BN的参数只有输入通道数，即通道数保持不变
elif m is Concat:
c2 = sum(ch[x] for x in f)Concat：f是所有需要拼接层的索引，则输出通道c2是所有层的和
elif m is Detect:
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)
#[list(range(args[1] * 2))]：初始化列表：预测框的宽高
#最后乘上len(f) ，就是生成所有预测层对应的预测框的初始高宽Detect类：
模块信息args先填入每个预测层的输入通道数，然后填入生成所有预测层对应的预测框的初始高宽的列表。
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
else:
c2 = ch[f]#其余情况都是输出通道数（c2）为输入通道数Contract和Expand类在网络结构中没有，咱就不说他了
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
#拿args里的参数去构建了module m，然后模块的循环次数用参数n控制。拿args里的参数去构建了module m，然后模块的循环次数用参数n控制
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
#以上是日志文件信息（每一层module构建的编号、参数量等）以上是日志文件信息（每一层module构建的编号、参数量等）
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
layers.append(m_)#把构建的模块保存到layers里
if i == 0:
ch = []
ch.append(c2)#把该层的输出通道数写入ch列表里保存需要用的层的输出（比如Concat层需要concat某些层，这些层的结果就需要存起来），这里后续会在_forward_once函数中用到
构建的模块保存在layers里，并且把该层的输出通道数写进ch中
#当循环结束后再构建成模型
return nn.Sequential(*layers), sorted(save)循环结束后再构建成模型。
五、代码分析注释全家桶
1、Model部分
class Model(nn.Module):
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
super().__init__()
if isinstance(cfg, dict): #判断cfg是不是dict（字典）类型
self.yaml = cfg # model dict
#模型词典赋值给self.yaml
else: # is *.yaml
import yaml # for torch hub
self.yaml_file = Path(cfg).name # 获取cfg的文件名
with open(cfg, encoding='ascii', errors='ignore') as f:#用ascii编码，忽略错误的形式打开文件cfg
self.yaml = yaml.safe_load(f) # model dict
#用yaml的文件形式加载cfg，赋值给self.yaml
# 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
#上面是为了判断输入的channel和配置文件里的是否相同，不相同则变成输入的参数。
if anchors:
LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
self.yaml['anchors'] = round(anchors) # override yaml value
#以上将anchors进行四舍五入（防止输入的是小数从而报错）
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
#给那些种类编编号，从0到nc-1
self.inplace = self.yaml.get('inplace', True)
# Build strides, anchors
m = self.model[-1] # Detect()
if isinstance(m, Detect):
s = 256 # 2x min stride步长初始化
m.inplace = self.inplace
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
m.anchors /= m.stride.view(-1, 1, 1)
check_anchor_order(m) #根据YOLOv5 Detect（）模块m的步幅顺序检查锚定顺序，必要时进行纠正
self.stride = m.stride
self._initialize_biases() # only run once
#将偏差初始化进Detect模块（没有偏差，因为初始类频率cf为None）
# 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
#默认情况下不增强2、_forward_once函数
def _forward_once(self, x, profile=False, visualize=False):
y, dt = [], [] # outputs
for m in self.model:
#m中的参数：m.f是从哪层开始,m.n是模块的默认深度,m.args是该层的参数（就是从yaml那来的）
if m.f != -1: # if not from previous layer
#不是上一层,就是说不直接连接上一层（比如像Concat)
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
#取出对应的层的结果，准备后面的进入对应m的forward()
if profile:
self._profile_one_layer(m, x, dt)
x = m(x) # run
#m是模块（某个层）的意思，所以x传入模块，相当于执行模块（比如说Focus，SPP等）中的forward（）
#第一层Focus的m.f是-1，所以直接跳到这一步开始
y.append(x if m.i in self.save else None) # save output
#将每一层的输出结果保存到y
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
return x3、parse_model函数
def parse_model(d, ch): # model_dict, input_channels(3)
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
#日志记载，不管他
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
#以上是读取配置dict里面的参数
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 # 将模块类型m转为值
for j, a in enumerate(args):#循环模块参数args
try:
args[j] = eval(a) if isinstance(a, str) else a # 将每个模块参数args[j]转为值
except NameError:
pass
#以上循环，开始迭代循环backbone与head的配置
#f：从哪层开始;n：模块的默认深度;m：模块的类型;args：模块的参数
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
#网络用（n*gd）控制模块的深度缩放。
#深度在这里指的是像CSP这种模块的重复迭代次数，宽度一般我们指的是特征图的channel。
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
c1, c2 = ch[f], args[0]
#ch是用来保存之前所有的模块输出的channle（所以ch[-1]代表着上一个模块的输出通道）。
# ch[f]是第f层的输出。args[0]是默认的输出通道。
if c2 != no: # if not output
#如果不是最终的输出
c2 = make_divisible(c2 * gw, 8)#保证了输出的通道是8的倍数
args = [c1, c2, *args[1:]] #args变为原来的args+module的输入通道数（c1）、输出通道数（c2）
if m in [BottleneckCSP, C3, C3TR, C3Ghost]: #只有CSP结构的才会根据深度参数n来调整该模块的重复迭加次数
args.insert(2, n) # number of repeats
#模块参数信息args插入n
n = 1#n重置
elif m is nn.BatchNorm2d:
args = [ch[f]]#BN的参数只有输入通道数，即通道数保持不变
elif m is Concat:
c2 = sum(ch[x] for x in f)#Concat：f是所有需要拼接层的索引，则输出通道数c2是所有层的和
elif m is Detect:
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)
#[list(range(args[1] * 2))]：初始化列表：预测框的宽高
#最后乘上len(f) ，就是生成所有预测层对应的预测框的初始高宽
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
else:
c2 = ch[f]#其余情况都是输出通道数（c2）为输入通道数
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
#拿args里的参数去构建了module m，然后模块的循环次数用参数n控制。整体都受到宽度缩放，C3模块受到深度缩放。
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
#以上是日志文件信息（每一层module构建的编号、参数量等）
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
#保存需要用的层的输出（比如Concat层需要concat某些层，这些层的结果就需要存起来）
layers.append(m_)#把构建的模块保存到layers里
if i == 0:
ch = []
ch.append(c2)#把该层的输出通道数写入ch列表里
#当循环结束后再构建成模型
return nn.Sequential(*layers), sorted(save)欢迎大家在评论区批评指正~
同时呢，原创不易，嘻嘻，如果喜欢的话，麻烦点点赞呗
