如何冻结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)

结论

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