020、迁移学习与预训练权重COCO 预训练的冻结策略与逐层解冻的最佳实践一个让我熬夜到凌晨三点的bug去年秋天我在给一个工业检测项目做YOLOv8的迁移学习。客户要求检测的是某种特殊钢材表面的微裂纹数据集只有800张标注图。我心想COCO预训练权重一加载fine-tune一下不就行了结果训练了50个epochmAP0.5只有0.12比随机初始化还差。当时我盯着tensorboard上那条几乎平躺的loss曲线脑子里只有一个念头冻结策略搞反了。后来我花了整整一周把YOLO的backbone、neck、head逐层拆开用hook打印了每一层的梯度变化才真正搞明白迁移学习在YOLO上应该怎么玩。今天这篇笔记就是那次踩坑的完整复盘。先搞清楚YOLO的预训练权重到底学了什么COCO数据集有80类从人、车到杯子、香蕉什么都有。YOLO在COCO上训练出来的权重backbone部分学到的是通用的视觉特征——边缘、纹理、形状、物体部件的组合模式。这部分对任何视觉任务都有用。但head部分检测头学的是COCO类别特有的分类和定位知识。比如“人的肩膀和头部的相对位置”、“车的前脸特征”。如果你要检测的是钢材裂纹这些知识不仅没用还会产生严重的负迁移。我见过有人直接把整个模型冻结只训练最后的分类层结果模型把裂纹误判成“领带”或者“香蕉”——因为COCO里没有裂纹模型只能从已有的80类里硬凑。核心原则backbone可以大胆复用head必须从头学起。冻结策略的三种典型场景场景一数据量极少500张这时候最怕过拟合。我的做法是# 加载预训练权重modelYOLO(yolov8n.pt)# 冻结backbone所有层forname,paraminmodel.model.model[:10].named_parameters():# YOLOv8 backbone共10层param.requires_gradFalse# 这里踩过坑只冻结了卷积层漏了BN层结果BN的running_mean还在更新# 别这样写if conv in name: param.requires_grad False注意YOLOv8的model.model是一个nn.Sequential前10层是backbone第10层到第21层是neck最后是head。不同版本索引可能不同建议先打印结构确认。冻结后只训练neck和head学习率设小一点比如1e-4。我一般先跑20个epoch看看loss有没有下降趋势没有的话说明冻结太狠了。场景二中等数据量500-3000张这是最常见的场景。我的策略是冻结backbone的前半部分解冻后半部分。为什么因为backbone浅层学的是边缘、颜色等低级特征这些在任何数据集上都是通用的。深层学的是物体部件组合跟具体任务相关度更高。# 冻结backbone前6层浅层特征fori,(name,param)inenumerate(model.model.model[:6].named_parameters()):param.requires_gradFalse# 打印确认print(f冻结: {name})# 后4层backbone 全部neck head 可训练# 注意这里有个坑BN层的weight和bias即使requires_gradFalse# 如果模型处于train模式BN的统计量依然会更新# 解决方案设置model.eval()模式但这样会影响dropout# 我一般手动设置BN的track_running_statsFalse这里有个细节BN层的处理。如果你冻结了某个卷积层最好连它后面的BN层一起冻结。否则BN的running_mean和running_var会持续更新导致特征分布偏移。我见过有人因为这个原因验证集loss震荡得像心电图。场景三数据量大3000张且领域差异大比如你要从COCO的通用物体检测迁移到卫星图像的目标检测。这时候COCO的预训练权重只能提供最底层的纹理特征高层特征基本没用。我的做法是只冻结backbone前3层其余全部解冻。学习率用warmup余弦退火初始学习率可以设到1e-3。# 只冻结最底层fori,(name,param)inenumerate(model.model.model[:3].named_parameters()):param.requires_gradFalse# 其余全部可训练但给不同层设置不同学习率# 这里用到了param_group的技巧optimizertorch.optim.SGD([{params:model.model.model[3:].parameters(),lr:1e-3},{params:model.model.model[:3].parameters(),lr:0},# 冻结层lr0],lr1e-3,momentum0.937,weight_decay5e-4)逐层解冻不是玄学是工程很多人问我“逐层解冻到底怎么操作是不是每几个epoch解冻一层”我的回答是看loss曲线不看epoch数。具体做法先冻结全部backbone只训练head。观察loss曲线等它进入平台期通常3-5个epoch。解冻backbone最后一层第10层继续训练。loss应该会继续下降。等loss再次进入平台期解冻第9层。以此类推直到所有层解冻。# 逐层解冻的伪代码defgradual_unfreeze(model,current_epoch,unfreeze_schedule): unfreeze_schedule: {5: 10, 10: 9, 15: 8, ...} 表示第5epoch解冻第10层第10epoch解冻第9层 ifcurrent_epochinunfreeze_schedule:layer_idxunfreeze_schedule[current_epoch]forname,paraminmodel.model.model[layer_idx-1].named_parameters():param.requires_gradTrue# 别这样写直接param.requires_grad True# 记得同时解冻对应的BN层# 我一般会打印print(f解冻第{layer_idx}层)注意解冻后学习率要适当降低我一般降为原来的0.1倍。因为新解冻的层参数还比较随机太大的学习率会把之前训练好的特征冲乱。一个容易忽略的细节预训练权重的类别数YOLOv8的head输出维度是(41num_classes)*3其中4是bbox坐标1是objectnessnum_classes是类别数。COCO是80类所以head输出是(4180)*3255。如果你要检测的类别数不是80直接加载预训练权重会报错。YOLO的源码里有个自动处理# 在YOLO的__init__里ifnc!self.nc:# 重新初始化head的最后一层卷积self.model[-1].cv2Conv(...)# bbox分支self.model[-1].cv3Conv(...)# cls分支但这里有个坑只重新初始化了最后一层卷积前面的层还是COCO的。这意味着如果你从80类改成1类head的前面几层依然保留着COCO的分类特征需要足够的数据来冲刷掉。我建议的做法是如果类别数变化较大比如从80类改成1类手动把整个head重新初始化不要依赖自动处理。# 手动重新初始化headdefreinit_head(model,num_classes):# 找到head部分headmodel.model.model[-1]# 重新初始化所有卷积层forminhead.modules():ifisinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,modefan_out,nonlinearityrelu)ifm.biasisnotNone:nn.init.constant_(m.bias,0)# 修改类别数head.ncnum_classes实战中的几个血泪教训教训一不要冻结BN层但让模型处于train模式BN层在train模式下会更新running_mean和running_var。如果你冻结了某个卷积层但没冻结BN层BN的统计量会持续变化导致特征分布偏移。解决方案要么把BN也冻结要么设置model.eval()。教训二学习率不是越小越好我见过有人把学习率设到1e-6训练了100个epoch loss纹丝不动。迁移学习的学习率一般比从头训练小10倍但也不能太小。我通常用1e-4作为初始值观察loss下降速度如果太慢就调大。教训三验证集和测试集要分开有一次我偷懒用验证集当测试集结果模型在验证集上mAP 0.85上线后只有0.3。后来发现是因为验证集和训练集来自同一个批次的数据分布太相似。迁移学习的数据集划分更要谨慎因为预训练模型可能已经见过类似的数据。我的个人经验总结先跑一个epoch看看loss加载预训练权重后先跑一个epoch看loss初始值。如果loss比随机初始化还高说明预训练权重不适合你的任务考虑换backbone或者从头训练。冻结策略不是一成不变的我一般会准备3-4种冻结方案每种跑5个epoch看哪个方案的loss下降最快。选最好的继续训练。不要迷信“全部解冻”很多人觉得全部解冻效果最好其实不然。对于小数据集全部解冻很容易过拟合。我见过一个案例冻结backbone比全部解冻mAP高了15个点。记录每次实验的冻结策略我习惯在实验名称里加上冻结信息比如“frozen_6layers_lr1e-4”。这样回头复盘时能快速定位问题。最后也是最重要的迁移学习不是银弹。如果你的数据集和COCO差异太大比如医学影像、卫星图像预训练权重的帮助有限。这时候考虑用ImageNet预训练的backbone或者干脆用自监督预训练。写这篇笔记的时候我翻出了去年那个钢材裂纹项目的实验记录。最终方案是冻结backbone前6层解冻后4层全部neckhead学习率1e-4余弦退火50个epoch。mAP0.5从0.12涨到了0.78。虽然不算惊艳但至少能用了。希望这篇笔记能帮你少踩几个坑。下次遇到迁移学习的问题记得先看看loss曲线再决定怎么冻结。