pytorch搭建YOLOv5网络
1.基本组成模块
1.1 Foucs模块
# --------------------------Foucs模块-------------——-
class Foucs(nn.Module):
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups =1普通卷积 >1深度可分离卷积
super(Foucs, self).__init__()
self.conv = Conv(c1*4, c2, k, s, p, g, act)
def forward(self, x):
# focus层：每隔一个像素取一个，然后进行堆叠，每个通道转变为四个特征层
# torch.cat((a,b,c),dim) dim为几按第几为拼接
# (batch,3,640,640）——》（batch,12,320,320）
img = torch.cat(
[
x[..., ::2, ::2],
x[..., 1::2, ::2],
x[..., ::2, 1::2],
x[..., 1::2, 1::2]
], 1
)
# 第一层卷积
img = self.conv(img)
return img1.2 Conv模块
# -------------------------Conv模块------------------
# Conv这个函数是整个网络中最基础的组件，由卷积层 + BN层 + 激活函数 组成
class Conv(nn.Module):
def __init__(self, c1, c2, k, s, p=None, g=1, act=True):
super(Conv, self).__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # 卷积层
self.bn = nn.BatchNorm2d(c2, eps=0.001, momentum=0.03) # 批归一化
# 激活函数类型,True就是SiLU() / Swish,False就是不使用激活函数,类型是nn.Module就使用传进来的激活函数类型
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)))
"""用于Model类的fuse函数
融合conv+bn 加速推理 一般用于测试/验证阶段
"""
def fuseforward(self, x):
return self.act(self.conv(x))
1.3 Bottlenet模块
# --------------------------Bottlenet模块-------------------
# 残差模块，类似于颈部，先减少通道数，再增加通道数，输入前后通道数不变，通过增加相当的深度来提高准确率。
class Bottleneck(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
"""
:param n: 有n个Bottleneck
:param shortcut: bool Bottleneck中是否有shortcut，默认True 是否有大残差边
:param g:=1普通卷积 >1深度可分离卷积
:param e: c2*e=中间其他所有层的卷积核个数（中间层的输入输出channel数）
"""
super(Bottleneck, self).__init__()
c_ = int(c2 * e) # 中间层通道数
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))1.4 C3模块
# -------------------------————C3模块-------------------------
# 卷积核全为1，宽高不改变
class C3(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
"""
:param n: 有n个Bottleneck
:param shortcut: bool Bottleneck中是否有shortcut，默认True
:param g:=1普通卷积 >1深度可分离卷积
:param e: c2*e=中间其他所有层的卷积核个数（中间层的输入输出channel数）
"""
super(C3, self).__init__()
c_ = int(c2 * e)
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for i in range(n)]) # *将列表拆成一个个元素输出
def forward(self, x):
return self.cv3(torch.cat
(
(self.cv1(x), self.m(self.cv2(x))), dim=1
))
1.5 SPP模块
# -------------------------————SPP模块-------------------------
# 将更多不同分辨率的特征进行融合,经过不同卷积核最大池化MaxPool2d后通道和大小都不变，然后进行通道叠加
class SPP(nn.Module):
def __init__(self, c1, c2, k_list=(5, 9, 13)):
'''
:param k_list: 三个卷积核尺寸
'''
super(SPP, self).__init__()
c_ = c1 // 2
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv((len(k_list) + 1) * c_, c2, 1, 1)
self.m = nn.ModuleList([nn.MaxPool2d(k, 1, k // 2) for k in k_list])
def forward(self, x):
x = self.cv1(x)
return self.cv2(torch.cat([x] + [m(x) for m in self.m], dim=1)) # [a]+[b,c]=[a,b,c]1.6 SPPF模块（效果和SPP一样，但用小卷积核代替大卷积核，计算量更小）
# -------------------------————SPPF模块-------------------------
# SPPF结构是将输入串行通过多个5x5大小的MaxPool层，这里需要注意的是串行两个5x5大小的MaxPool层是和一个9x9大小的MaxPool层计算结果是一样的，
# 串行三个5x5大小的MaxPool层是和一个13x13大小的MaxPool层计算结果是一样的。
class SPPF(nn.Module):
def __init__(self, c1, c2, k=5):
super(SPPF, self).__init__()
c_ = c1 // 2
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(4 * c_, c2, 1, 1)
self.m = nn.MaxPool2d(k, 1, k // 2)
def forward(self, x):
x = self.cv1(x)
x1 = self.m(x)
x2 = self.m(x1)
x3 = self.m(x2)
return self.cv2(torch.cat([x, x1, x2, x3], dim=1))
2 主干网络搭建
# -----------------------主干网络CSPdarknet-------------------
class CSPdarknet(nn.Module):
def __init__(self, base_channel, n, phi, pretrain):
"""
:param base_channel:初始基本通道数，经过Fouse+Conv
:param n:有n个Bottleneck:
:param index: 预训练模型索引“s m l x”
:param pretrain: 是否使用预训练模型
"""
super(CSPdarknet, self).__init__()
# ————————------———第一层———————————--————#
'''
输入图片3 640 640 初始基本通道base_channel是64
Fouse+Conv
3 640 640 -> 12 320 320 -> 64 320 320
'''
self.fouse = Foucs(c1=3, c2=base_channel, k=3)
# ————————------———第二层———————————--————#
"""
Conv + C3
64 320 320 -> 128 160 160 -> 128 160 160
"""
self.dark1 = nn.Sequential(
Conv(base_channel, base_channel * 2, 3, 2),
C3(base_channel * 2, base_channel * 2, n)
)
# ————————------———第三层———————————--————#
"""
Conv + C3
128 160 160 -> 256 80 80 ->256 80 80
完成此步之后，引出第1个有效特征层256 80 80
"""
self.dark2 = nn.Sequential(
Conv(base_channel * 2, base_channel * 4, 3, 2),
C3(base_channel * 4, base_channel * 4, n * 3)
)
# ————————------———第四层———————————--————#
"""
Conv + C3
256 80 80 -> 512 40 40 ->512 40 40
完成此步之后，引出第2个有效特征层512 40 40
"""
self.dark3 = nn.Sequential(
Conv(base_channel * 4, base_channel * 8, 3, 2),
C3(base_channel * 8, base_channel * 8, n * 3)
)
# ————————------———第五层———————————--————#
"""
Conv + SPPF +C3
512 40 40 -> 1024 20 20 -> 1024 20 20 -> 1024 20 20
完成此步之后，引出第3个有效特征层1024 20 20
"""
self.dark4 = nn.Sequential(
Conv(base_channel * 8, base_channel * 16, 3, 2),
SPPF(base_channel * 16, base_channel * 16),
C3(base_channel * 16, base_channel * 16, n, shortcut=False)
)
def forward(self, x):
x = self.fouse(x)
x = self.dark1(x)
feat1 = self.dark2(x)
feat2 = self.dark3(feat1)
feat3 = self.dark4(feat2)
return feat1, feat2, feat3 # 三个不同尺度的有效特征层