066、优化器参数分组策略：Weight Decay 只作用于 Weight 不作用于 Bias 的实现

066、优化器参数分组策略Weight Decay 只作用于 Weight 不作用于 Bias 的实现一个让我debug到凌晨三点的bug去年做YOLOv5的蒸馏实验模型在COCO上训到第80个epoch突然loss炸了。检查了学习率、数据增强、BN层都没问题。最后发现是优化器参数分组写错了——Weight Decay打到了Bias上导致某些层的偏置项在训练后期被过度正则化梯度直接崩掉。这个坑我踩了整整两天。后来翻YOLOv5源码发现Ultralytics团队在优化器初始化时专门做了参数分组把Weight Decay只施加在Weight上Bias和BN层的参数全部豁免。今天就把这个细节拆开揉碎了讲清楚。为什么Weight Decay不能碰Bias先明确一个基本认知Weight Decay的本质是在损失函数中增加L2正则项让权重向零收缩。但Bias的作用是提供平移不变性如果对Bias施加正则化相当于强制模型把决策边界往原点拉这会破坏特征的偏移补偿能力。举个例子BN层有γ和β两个可学习参数γ负责缩放β负责平移。如果对β施加Weight Decay模型会倾向于把β推向0导致BN层的平移能力被削弱训练后期特征分布偏移时无法有效补偿。更直观的全连接层的Bias如果被正则化相当于在分类时强行让所有类别的决策边界都经过原点这在多分类任务中几乎不可能达到最优。PyTorch优化器的参数分组机制PyTorch的优化器支持传入参数组列表每组可以独立设置学习率、权重衰减等超参数。标准写法是这样的optimizertorch.optim.SGD([{params:model.backbone.parameters(),lr:0.01},{params:model.head.parameters(),lr:0.001}],weight_decay0.0005)但这里有个陷阱如果只传一个参数组所有参数共享同一个weight_decay。而我们需要的是——Weight只加weight_decayBias和BN参数不加。YOLOv5源码中的实现逐行解析直接看YOLOv5的train.py里优化器初始化部分我把它简化成可复用的函数defgroup_parameters(model,weight_decay0.0005): 参数分组weight_decay只作用于weight不作用于bias和BN参数 这里踩过坑如果直接对model.parameters()统一设置weight_decaybias会被正则化 # 分别收集需要和不需要weight_decay的参数decay_params[]# 需要weight_decay的参数主要是weightno_decay_params[]# 不需要weight_decay的参数bias和BN参数forname,paraminmodel.named_parameters():ifnotparam.requires_grad:continue# 冻结的参数直接跳过别浪费计算# 判断条件如果参数名包含bias或者属于BN层weight和bias都豁免# 别这样写if bias in name or bn in name —— 会漏掉BN的weightifbiasinnameorbninnameornorminname:no_decay_params.append(param)else:decay_params.append(param)# 构建参数组注意weight_decay只在decay组设置optimizertorch.optim.SGD([{params:decay_params,weight_decay:weight_decay},{params:no_decay_params,weight_decay:0.0}],lr0.01,momentum0.937)# 这里lr和momentum是全局默认值returnoptimizer关键点在于weight_decay参数在参数组级别设置全局的weight_decay会被参数组内的设置覆盖。所以decay组显式设置weight_decay0.0005no_decay组设置weight_decay0.0。一个容易忽略的细节BN层的weight也需要豁免很多人以为BN层只有bias需要豁免其实BN层的weightγ参数同样不应该加weight_decay。原因在于BN的γ是控制特征缩放幅度的如果被正则化模型会倾向于把γ推向1因为L2正则会让参数变小这会破坏BN的自适应缩放能力。YOLOv5源码中判断条件写的是bias in name or bn in name这个bn in name会同时匹配BN层的weight和bias。如果你的模型里BN层命名不是标准的bn开头比如用了BatchNorm2d的默认命名记得检查一下参数名。进阶针对不同层设置不同的weight_decay有时候我们需要更精细的控制比如对backbone和head设置不同的weight_decaydefadvanced_group_parameters(model,backbone_wd0.0005,head_wd0.001): 进阶版backbone和head使用不同的weight_decay head层通常需要更强的正则化防止过拟合 decay_params{backbone:[],head:[]}no_decay_params{backbone:[],head:[]}forname,paraminmodel.named_parameters():ifnotparam.requires_grad:continue# 判断属于backbone还是head这里假设模型有model.backbone和model.headifbackboneinname:partbackboneelifheadinname:partheadelse:continue# 其他部分按默认处理ifbiasinnameorbninnameornorminname:no_decay_params[part].append(param)else:decay_params[part].append(param)optimizertorch.optim.SGD([{params:decay_params[backbone],weight_decay:backbone_wd},{params:no_decay_params[backbone],weight_decay:0.0},{params:decay_params[head],weight_decay:head_wd},{params:no_decay_params[head],weight_decay:0.0}],lr0.01,momentum0.937)returnoptimizer这个写法在迁移学习场景下特别有用——冻结backbone时只需要把backbone的参数组去掉即可。验证分组是否生效写完了分组逻辑怎么确认真的生效了别靠感觉直接打印参数组信息definspect_optimizer(optimizer):检查优化器参数分组是否正确调试用fori,param_groupinenumerate(optimizer.param_groups):wdparam_group.get(weight_decay,0)lrparam_group.get(lr,0)num_paramslen(param_group[params])# 取第一个参数的名字做示例需要提前保存参数名到param_groupprint(fGroup{i}: lr{lr}, weight_decay{wd}, params_count{num_params})建议在训练脚本里加一个--debug_optim参数开启后打印所有参数组信息确保bias和BN参数确实在weight_decay0的组里。一个常见的坑混合精度训练时的参数分组如果你用了AMP自动混合精度注意torch.cuda.amp.GradScaler不会影响参数分组但优化器状态如Adam的动量会随着参数分组独立维护。所以分组逻辑在AMP下同样适用不需要额外修改。但有一个细节如果用了torch.compile编译模型model.named_parameters()返回的参数名可能会被修改比如加上_orig_mod前缀这时候判断条件里的backbone in name可能失效。解决方案是在编译前保存参数名映射或者直接用model.backbone.named_parameters()这种子模块遍历方式。个人经验建议不要偷懒用model.parameters()统一设置weight_decay除非你的模型只有卷积层没有BN层这种情况几乎不存在。我见过太多人因为这个导致训练不稳定还以为是学习率的问题。BN层的γ和β都要豁免别只豁免bias。YOLOv5源码里用bn in name一次性搞定这个写法值得借鉴。如果用了自定义层比如SE模块、注意力机制注意检查这些层里有没有bias参数。有些实现会在全连接层加bias这些bias同样需要豁免。调试时打印参数组信息不要相信直觉。我曾经以为分组写对了结果打印出来发现所有参数都在同一个组里——因为参数名判断条件写错了。迁移学习时特别注意冻结的backbone参数如果还在优化器参数组里虽然梯度为0不会更新但weight_decay仍然会作用因为weight_decay是在梯度更新后直接对参数做衰减。所以冻结的参数最好从优化器参数组中移除或者设置weight_decay0。最后说一句参数分组这个细节看起来只是几行代码的事但搞不好能让你的模型训练直接崩掉。YOLOv5能稳定训练几百个epoch这种细节功不可没。下次写优化器初始化的时候多花两分钟把分组逻辑写对能省下后面debug的无数时间。