[YOLOv5-6.x] Set learnable weights combined with BiFPN

[YOLOv5-6.x] Set learnable weights combined with BiFPN

Foreword

In the previous blog, I briefly introduced the principle of BiFPN and how the author of YOLOv5 combined BiFPN: [Magic modification of YOLOv5-6.x (middle)]: adding ACON activation function, CBAM and CA attention mechanism, weighted bidirectional features Pyramid BiFPN

This article will try to further combine BiFPN, mainly from: YOLOv5 combined with BiFPN

Modify the yaml file (take yolov5s as an example)
only one modification
This article takes yolov5s.yaml as an example to modify. When modifying the model configuration file, pay attention to the following points:

The yaml file here is only modified in one place, that is, the Concat of the 19th layer is replaced by BiFPN_Add. If you want to modify the Concat of other layers, you can modify it by analogy.
The essence of BiFPN_Add is an add operation, not a concat operation. Therefore, each input layer of BiFPN_Add requires exactly the same size (number of channels, feature map size, etc.), so here we need to modify the previous parameters [-1, 13, 6] to meet This request:
-1 layer is the output of the previous layer. The original output channel number of the previous layer is 256, which is changed to 512 here.
The 13th floor is here [-1, 3, C3, [512, False]], # 13
After this modification, each input size of BiFPN_Add is [bs,256,40,40]
Finally, the parameter layer behind BiFPN_Add is set to [256, 256], that is, the number of input and output channels is 256

# YOLOv5 ? by Ultralytics, GPL-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 BiFPN 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, [512, 3, 2]], # 为了BiFPN正确add,调整channel数
[[-1, 13, 6], 1, BiFPN_Add3, [256, 256]], # cat P4 <— BiFPN change 注意v5s通道数是默认参数的一半
[-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)
]

Replace Concat all with BiFPN_Add

# YOLOv5 ? by Ultralytics, GPL-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 BiFPN head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, ‘nearest’]],
[[-1, 6], 1, BiFPN_Add2, [256, 256]], # 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, BiFPN_Add2, [128, 128]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)

[-1, 1, Conv, [512, 3, 2]], # 为了BiFPN正确add,调整channel数
[[-1, 13, 6], 1, BiFPN_Add3, [256, 256]], # cat P4 <— BiFPN change 注意v5s通道数是默认参数的一半
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)

[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, BiFPN_Add2, [256, 256]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)

[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]

print model parameters
You can refer to this blog: [YOLOv5-6.x] model parameters and detect layer output test (for personal use), test the model configuration file and view the output results:

from n params module arguments
0 -1 1 3520 models.common.Conv [3, 32, 6, 2, 2]
1 -1 1 18560 models.common.Conv [32, 64, 3, 2]
2 -1 1 18816 models.common.C3 [64, 64, 1]
3 -1 1 73984 models.common.Conv [64, 128, 3, 2]
4 -1 2 115712 models.common.C3 [128, 128, 2]
5 -1 1 295424 models.common.Conv [128, 256, 3, 2]
6 -1 3 625152 models.common.C3 [256, 256, 3]
7 -1 1 1180672 models.common.Conv [256, 512, 3, 2]
8 -1 1 1182720 models.common.C3 [512, 512, 1]
9 -1 1 656896 models.common.SPPF [512, 512, 5]
10 -1 1 131584 models.common.Conv [512, 256, 1, 1]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, ‘nearest’]
12 [-1, 6] 1 65794 models.common.BiFPN_Add2 [256, 256]
13 -1 1 296448 models.common.C3 [256, 256, 1, False]
14 -1 1 33024 models.common.Conv [256, 128, 1, 1]
15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, ‘nearest’]
16 [-1, 4] 1 16514 models.common.BiFPN_Add2 [128, 128]
17 -1 1 74496 models.common.C3 [128, 128, 1, False]
18 -1 1 295424 models.common.Conv [128, 256, 3, 2]
19 [-1, 13, 6] 1 65795 models.common.BiFPN_Add3 [256, 256]
20 -1 1 296448 models.common.C3 [256, 256, 1, False]
21 -1 1 590336 models.common.Conv [256, 256, 3, 2]
22 [-1, 10] 1 65794 models.common.BiFPN_Add2 [256, 256]
23 -1 1 1051648 models.common.C3 [256, 512, 1, False]
24 [17, 20, 23] 1 229245 models.yolo.Detect [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
Model Summary: 278 layers, 7384006 parameters, 7384006 gradients, 17.2 GFLOPs

Modify common.py
Copy and paste the code:

# 结合BiFPN 设置可学习参数 学习不同分支的权重
# 两个分支add操作
class BiFPN_Add2(nn.Module):
def __init__(self, c1, c2):
super(BiFPN_Add2, self).__init__()
# 设置可学习参数 nn.Parameter的作用是:将一个不可训练的类型Tensor转换成可以训练的类型parameter
# 并且会向宿主模型注册该参数 成为其一部分 即model.parameters()会包含这个parameter
# 从而在参数优化的时候可以自动一起优化
self.w = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
self.epsilon = 0.0001
self.conv = nn.Conv2d(c1, c2, kernel_size=1, stride=1, padding=0)
self.silu = nn.SiLU()

def forward(self, x):
w = self.w
weight = w / (torch.sum(w, dim=0) + self.epsilon)
return self.conv(self.silu(weight[0] * x[0] + weight[1] * x[1]))

# 三个分支add操作
class BiFPN_Add3(nn.Module):
def __init__(self, c1, c2):
super(BiFPN_Add3, self).__init__()
self.w = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
self.epsilon = 0.0001
self.conv = nn.Conv2d(c1, c2, kernel_size=1, stride=1, padding=0)
self.silu = nn.SiLU()

def forward(self, x):
w = self.w
weight = w / (torch.sum(w, dim=0) + self.epsilon) # 将权重进行归一化
# Fast normalized fusion
return self.conv(self.silu(weight[0] * x[0] + weight[1] * x[1] + weight[2] * x[2]))

Modify yolo.py
Find the elif m is Concat: statement in the parse_model function, and add the BiFPN_Add related statement after it:

elif m is Concat:
c2 = sum(ch[x] for x in f)
# 添加bifpn_add结构
elif m in [BiFPN_Add2, BiFPN_Add3]:
c2 = max([ch[x] for x in f])

Modify train.py
1. Add the weight parameters of BiFPN to the optimizer
Add the w parameter defined in the BiFPN_Add2 and BiFPN_Add3 functions to g1

g0, g1, g2 = [], [], [] # optimizer parameter groups
for v in model.modules():
# hasattr: 测试指定的对象是否具有给定的属性,返回一个布尔值
if hasattr(v, ‘bias’) and isinstance(v.bias, nn.Parameter): # bias
g2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d): # weight (no decay)
g0.append(v.weight)
elif hasattr(v, ‘weight’) and isinstance(v.weight, nn.Parameter): # weight (with decay)
g1.append(v.weight)
# BiFPN_Concat
elif isinstance(v, BiFPN_Add2) and hasattr(v, ‘w’) and isinstance(v.w, nn.Parameter):
g1.append(v.w)
elif isinstance(v, BiFPN_Add3) and hasattr(v, ‘w’) and isinstance(v.w, nn.Parameter):
g1.append(v.w)

2. Check the update status of BiFPN_Add layer parameters
If you want to check the parameter update status of the BiFPN_Add layer, you can refer to this blog [Pytorch] to view the parameter value of a certain layer of the model (for your own use), directly locate the w parameter, and output the corresponding value as the model is trained.