前言
上一篇中我們簡單回顧了yolov5的架構和搭建了開發環境,需要回顧的小伙伴可以點擊下面傳送門,基于第一篇的理解我們再進行原始碼解讀就會有事半功倍的效果,

yolov5架構剖析和環境搭建傳送門
這篇中我們開始決議yolov5的原始碼,類似yolov3我們同樣從模型的構建開始,這部分核心的代碼在yolo.py檔案中,我采用的yolo代碼版本是官方4.0版本(2021年1月發布的最新版本),這個版本的很多細節都和之前的略有區別了,代碼傳送門見下面鏈接,
代碼傳送門
另外大家可以在官網看到yolov5的技術演進路線,看到每個版本的更新的內容:

在決議前仍然需要說的是理解原始碼最好的辦法就是配置好圖片和標簽,然后進行debug,官方提供了coco128的資料集可以快速下載跑通,方便debug,
我們列印出每行原始碼涉及的tensor的維度,標注好函式的功能,下面我們開始吧!
yolo.py 原始碼決議
1.init()函式
我們首先來看模型構建的原始碼部分:class Model()
首先是該類的init()函式的原始碼,模型在init()就已經通過組態檔全部構建了:
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes
super(Model, self).__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) as f:
self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict
# Define model
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
if nc and nc != self.yaml['nc']:
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
self.yaml['nc'] = nc # 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
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
# Build strides, anchors
m = self.model[-1] # Detect()
if isinstance(m, Detect):
s = 256 # 2x min stride
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)
self.stride = m.stride
self._initialize_biases() # only run once
# print('Strides: %s' % m.stride.tolist())
# Init weights, biases
initialize_weights(self)
self.info()
logger.info('')
前面的不用多說,模型仍然是決議yaml的組態檔:
我們以yolov5s為例進行debug:
with open(cfg) as f:
self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict

可以看到yaml決議后是一個dict形式:
- nc代表類別數量
- depth_multiple是控制模型深度的引數,
- width_multiple是一個控制模型寬度的引數,
- anchors是預置的錨框,FPN每層設定3個,共有3*3=9個,
- backbone是backbone網路的構建引數,根據這個配置可以加載出backbone網路,
- head是yolo head網路的構建引數,根據這個配置可以加載出yolo head的網路,(其實可以認為這部分是neck+head)
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
if nc and nc != self.yaml['nc']:
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
self.yaml['nc'] = nc # override yaml value
這里判斷一下輸入的channel和組態檔里的是否一致,不一致則以輸入引數為準,
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])
下面就要進入核心函式parse_model()了,這個函式的原始碼見下面:
ef parse_model(d, ch): # model_dict, input_channels(3)
logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
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)
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):
try:
args[j] = eval(a) if isinstance(a, str) else a # eval strings
except:
pass
n = max(round(n * gd), 1) if n > 1 else n # depth gain
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
C3]:
c1, c2 = ch[f], args[0]
if c2 != no: # if not output
c2 = make_divisible(c2 * gw, 8)
args = [c1, c2, *args[1:]]
if m in [BottleneckCSP, C3]:
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])
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)
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
else:
c2 = ch[f]
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('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # 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)
下面我們來逐步決議下這個函式:
logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
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里面的引數,na是判斷anchor的數量,no是根據anchor數量推斷的輸出維度,比如對于coco是255,輸出維度=anchor數量*(類別數量+置信度+xywh四個回歸坐標),
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):
try:
args[j] = eval(a) if isinstance(a, str) else a # eval strings
except:
pass
n = max(round(n * gd), 1) if n > 1 else n # depth gain
這里開始迭代回圈backbone與head的配置,f,n,m,args分別代表著從哪層開始,模塊的默認深度,模塊的型別和模塊的引數,
n = max(round(n * gd), 1) if n > 1 else n
網路用n*gd控制模塊的深度縮放,比如對于yolo5s來講,gd為0.33,也就是把默認的深度縮放為原來的1/3,
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
C3]:
c1, c2 = ch[f], args[0]
if c2 != no: # if not output
c2 = make_divisible(c2 * gw, 8)
args = [c1, c2, *args[1:]]
if m in [BottleneckCSP, C3]:
args.insert(2, n) # number of repeats
n = 1
對于以上的這幾種型別的模塊,ch是一個用來保存之前所有的模塊輸出的channle,ch[-1]代表著上一個模塊的輸出通道,args[0]是默認的輸出通道,
def make_divisible(x, divisor):
# Returns x evenly divisible by divisor
return math.ceil(x / divisor) * divisor
這里配合make_divisible()函式,是為了放縮網路模塊的寬度(既輸出的通道數),比如對于第一個模塊“Focus”,默認的輸出通道是64,而yolov5s里的放縮系數是0.5,所以通過以上代碼變換,最終的輸出通道為32,make_divisible()函式保證了輸出的通道是8的倍數,
args = [c1, c2, *args[1:]]
if m in [BottleneckCSP, C3]:
args.insert(2, n) # number of repeats
n = 1
經過以上處理,args里面保存的前兩個引數就是module的輸入通道數、輸出通道數,如果是BottleneckCSP或者是C3則還需要重新添加重復次數的引數n,
elif m is nn.BatchNorm2d:
args = [ch[f]]
elif m is Concat:
c2 = sum([ch[x] for x in f])
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)
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
else:
c2 = ch[f]
以上是其他幾種型別的Module,
如果是nn.BatchNorm2d則通道數保持不變,
如果是Concat則f是所有需要拼接層的index,則輸出通道c2是所有層的和,
如果是Detect則對應檢測頭,這部分后面再詳細講,
Contract和Expand目前未在模型中使用,
m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module
這里把args里的引數用于構建了module m,然后模塊的回圈次數用引數n控制,值得注意的是C3模塊被控制為n=1也就是外部構建只回圈一次,
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('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
這里做了一些輸出列印,可以看到每一層module構建的編號、引數量等情況,比如:
from n params module arguments
0 -1 1 3520 models.common.Focus [3, 32, 3]
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串列里,
待全部回圈結束后再構建成模型,至此模型就全部構建完畢了,
現在我們再回到yolo.py里剛剛呼叫parse_model的位置然后繼續完成init()函式的學習,
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
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
# Build strides, anchors
m = self.model[-1] # Detect()
if isinstance(m, Detect):
s = 256 # 2x min stride
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)
self.stride = m.stride
self._initialize_biases() # only run once
# print('Strides: %s' % m.stride.tolist())
# Init weights, biases
initialize_weights(self)
self.info()
logger.info('')
這里通過呼叫一次forward()函式,輸入了一個[1, C, 256, 256]的tensor,然后得到FPN輸出結果的維度,然后求出了下采樣的倍數stride:8,16,32,
最后把anchor除以以上的數值,將anchor放縮到了3個不同的尺度上,anchor的最終shape是[3,3,2],
至此init()函式已經完整的過了一遍,
2.各種Modules的原始碼決議
在網路構建的程序中涉及到了多種Modules,這些Modules默認在models檔案夾下面的common.py檔案里我們下面還過一下這些函式,
(1) 普通卷積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(Conv, self).__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
self.bn = nn.BatchNorm2d(c2)
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 fuseforward(self, x):
return self.act(self.conv(x))
普通的卷積,這里呼叫了autopad()函式計算了same-padding所需要的padding數量,
默認的激活函式是SiLU(),
SiLU函式形式:f(x)=x?σ(x)
導函式形式: f(x)=f(x)+σ(x)(1?f(x))

(2) 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(Bottleneck, self).__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c2, 3, 1, g=g)
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))
可以看出BottleNeck結構默認是先1x1卷積縮小channel為原來的1/2,再通過3x3卷積提取特征,如果輸入通道c1和3x3卷積輸出通道c2相等,則進行殘差輸出,shortcut引數控制是否進行殘差連接,
(3) BottleNeckCSP和C3
class BottleneckCSP(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(BottleneckCSP, self).__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
self.cv4 = Conv(2 * c_, c2, 1, 1)
self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
self.act = nn.LeakyReLU(0.1, inplace=True)
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
y1 = self.cv3(self.m(self.cv1(x)))
y2 = self.cv2(x)
return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
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(C3, self).__init__()
c_ = int(c2 * e) # hidden channels
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))
在common.py里實作了兩種csp結構:

BottleneckCSP就完全對應著上面的結構,但是作者在yoloV5 4.0的版本中將這部分結構改成了C3,C3的結構如下圖:

殘差之后的Conv被去掉了,激活函式從上面的LeakyRelu變為了SiLU,
(4) SPP
class SPP(nn.Module):
# Spatial pyramid pooling layer used in YOLOv3-SPP
def __init__(self, c1, c2, k=(5, 9, 13)):
super(SPP, self).__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
def forward(self, x):
x = self.cv1(x)
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
SPP模塊將輸入通道減半,然后分別做kernel size為5,9,13的maxpooling,最后將結過拼接,包含原始輸入的四組結果合并后通道應該是原來的2倍,
(5) Focus
class Focus(nn.Module):
# Focus wh information into c-space
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super(Focus, self).__init__()
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
# self.contract = Contract(gain=2)
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
# return self.conv(self.contract(x))
把feature map 切成四等分,然后疊加起來,最后的結果是通道數變為原來的四倍,resolution為原來的1/4(H,W分別減半),最后通過一個卷積調整通道數為預先設定,
以上我們完成了模型構建部分的部分內容,剩下的Model中的forward()前向傳播和YoloV5的detection head部分我們將在下一篇里面講述,
如果你覺得內容有用的話,請點贊關注支持下博主,有問題歡迎留言謝謝!
轉載請註明出處,本文鏈接:https://www.uj5u.com/houduan/265104.html
標籤:python
