对大佬的改进进行学习过程中发现的问题及解决,原文
yolov5增加AFPN-全新特征融合模块AFPN,效果完胜PAFPN_athrunsunny的博客-CSDN博客
1.common.py加入 导入 import torch.nn.functional as F
class Upsample(nn.Module):
def __init__(self, in_channels, out_channels, scale_factor=2):
super(Upsample, self).__init__()
self.upsample = nn.Sequential(
Conv(in_channels, out_channels, 1),
nn.Upsample(scale_factor=scale_factor, mode='bilinear')
)
# carafe
# from mmcv.ops import CARAFEPack
# self.upsample = nn.Sequential(
# BasicConv(in_channels, out_channels, 1),
# CARAFEPack(out_channels, scale_factor=scale_factor)
# )
def forward(self, x):
x = self.upsample(x)
return x
class Downsample(nn.Module):
def __init__(self, in_channels, out_channels,scale_factor=2):
super(Downsample, self).__init__()
self.downsample = nn.Sequential(
Conv(in_channels, out_channels, scale_factor, scale_factor, 0)
)
def forward(self, x):
x = self.downsample(x)
return x
class ASFF_2(nn.Module):
def __init__(self, inter_dim=512,level=0,channel=[64,128]):
super(ASFF_2, self).__init__()
self.inter_dim = inter_dim
compress_c = 8
self.weight_level_1 = Conv(self.inter_dim, compress_c, 1, 1)
self.weight_level_2 = Conv(self.inter_dim, compress_c, 1, 1)
self.weight_levels = nn.Conv2d(compress_c * 2, 2, kernel_size=1, stride=1, padding=0)
self.conv = Conv(self.inter_dim, self.inter_dim, 3, 1)
self.upsample = Upsample(channel[1],channel[0])
self.downsample = Downsample(channel[0],channel[1])
self.level = level
def forward(self, x):
input1, input2 = x
if self.level == 0:
input2 = self.upsample(input2)
elif self.level == 1:
input1 = self.downsample(input1)
level_1_weight_v = self.weight_level_1(input1)
level_2_weight_v = self.weight_level_2(input2)
levels_weight_v = torch.cat((level_1_weight_v, level_2_weight_v), 1)
levels_weight = self.weight_levels(levels_weight_v)
levels_weight = F.softmax(levels_weight, dim=1)
fused_out_reduced = input1 * levels_weight[:, 0:1, :, :] + \
input2 * levels_weight[:, 1:2, :, :]
out = self.conv(fused_out_reduced)
return out
class ASFF_3(nn.Module):
def __init__(self, inter_dim=512,level=0,channel=[64,128,256]):
super(ASFF_3, self).__init__()
self.inter_dim = inter_dim
compress_c = 8
self.weight_level_1 = Conv(self.inter_dim, compress_c, 1, 1)
self.weight_level_2 = Conv(self.inter_dim, compress_c, 1, 1)
self.weight_level_3 = Conv(self.inter_dim, compress_c, 1, 1)
self.weight_levels = nn.Conv2d(compress_c * 3, 3, kernel_size=1, stride=1, padding=0)
self.conv = Conv(self.inter_dim, self.inter_dim, 3, 1)
self.level = level
if self.level == 0:
self.upsample4x = Upsample(channel[2],channel[0], scale_factor=4)
self.upsample2x = Upsample(channel[1], channel[0], scale_factor=2)
elif self.level == 1:
self.upsample2x1 = Upsample(channel[2], channel[1], scale_factor=2)
self.downsample2x1 = Downsample(channel[0],channel[1], scale_factor=2)
elif self.level == 2:
self.downsample2x = Downsample(channel[1], channel[2], scale_factor=2)
self.downsample4x = Downsample(channel[0], channel[2], scale_factor=4)
def forward(self, x):
input1, input2, input3 = x
if self.level == 0:
input2 = self.upsample2x(input2)
input3= self.upsample4x(input3)
elif self.level == 1:
input3 = self.upsample2x1(input3)
input1 = self.downsample2x1(input1)
elif self.level == 2:
input1 = self.downsample4x(input1)
input2 = self.downsample2x(input2)
level_1_weight_v = self.weight_level_1(input1)
level_2_weight_v = self.weight_level_2(input2)
level_3_weight_v = self.weight_level_3(input3)
levels_weight_v = torch.cat((level_1_weight_v, level_2_weight_v, level_3_weight_v), 1)
levels_weight = self.weight_levels(levels_weight_v)
levels_weight = F.softmax(levels_weight, dim=1)
fused_out_reduced = input1 * levels_weight[:, 0:1, :, :] + \
input2 * levels_weight[:, 1:2, :, :] + \
input3 * levels_weight[:, 2:, :, :]
out = self.conv(fused_out_reduced)
return out
2.yolov5s.yaml改为
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
# 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
# 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
]
# YOLOv5 v6.0 head
head:
[[4, 1, Conv, [128, 1, 1]], # 10
[6, 1, Conv, [256, 1, 1]], # 11
[[10, 11], 1, ASFF_2, [128, 0]], # 12
[[10, 11], 1, ASFF_2, [256, 1]], # 13
[-2, 1, C3, [128, False]], # 14
[-2, 1, C3, [256, False]], # 15
[9, 1, Conv, [512, 1, 1]], # 16
[[14, 15, 16], 1, ASFF_3, [128, 0]], # 17
[[14, 15, 16], 1, ASFF_3, [256, 1]], # 18
[[14, 15, 16], 1, ASFF_3, [512, 2]], # 19
[17, 1, C3, [128, False]],
[18, 1, C3, [256, False]],
[19, 1, C3, [512, False]],
[[20, 21, 22], 1, Detect, [nc, anchors]]
]
3.yolo.py 加入
elif m in {ASFF_2, ASFF_3}:
c2 = args[0]
if c2 != no: # if not output
c2 = make_divisible(c2 * gw, 8)
args[0] = c2
args.append([ch[x] for x in f])
位置在如图,def parse_model(d, ch): # model_dict, input_channels(3)下
4.train.py backward 附件加入
torch.use_deterministic_algorithms(False)
如图
总体来说和BiFPNl类似