yolov5 冻结网络

参考github上的指导

1、找到train.py中的freeze

 # Freeze 
 freeze = []  # parameter names to freeze (full or partial) 
 for k, v in model.named_parameters(): 
     v.requires_grad = True  # train all layers 
     if any(x in k for x in freeze): 
         print('freezing %s' % k) 
         v.requires_grad = False 

2、查看model的module names

for k, v in model.named_parameters():
    print(k)

# Output
model.0.conv.conv.weight
model.0.conv.bn.weight
model.0.conv.bn.bias
model.1.conv.weight
model.1.bn.weight
model.1.bn.bias
model.2.cv1.conv.weight
model.2.cv1.bn.weight
...
model.23.m.0.cv2.bn.weight
model.23.m.0.cv2.bn.bias
model.24.m.0.weight
model.24.m.0.bias
model.24.m.1.weight
model.24.m.1.bias
model.24.m.2.weight
model.24.m.2.bias

3、查看yolov5s.yaml
可以看到，0~9层是backbone

YOLOv5 backbone

backbone:

[from, number, module, args]

[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, BottleneckCSP, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, BottleneckCSP, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, BottleneckCSP, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, BottleneckCSP, [1024, False]], # 9
]

YOLOv5 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, BottleneckCSP, [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, BottleneckCSP, [256, False]],  # 17 (P3/8-small) 

[-1, 1, Conv, [256, 3, 2]], 
[[-1, 14], 1, Concat, [1]],  # cat head P4 
[-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium) 

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

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

]

4、冻结backbone
冻结’model.0.’ 到 ‘model.9.’ (注意这里两个‘.’都不能少)

freeze = ['model.%s.' % x for x in range(10)]  # parameter names to freeze (full or partial)

结论
冻结可以大幅加快训练速度且节省GPU显存，但会稍微降低最终精度。
在coco预训练过的yolov5m，在voc上冻结backbone训练约70个epochs，最终AP50可达89.03%，比不冻结低了0.21%