Fine-tuning 学习率策略：分层学习率衰减（Layer-wise Learning Rate Decay, LLRD）

Fine-tuning 学习率策略分层学习率衰减Layer-wise Learning Rate Decay, LLRDflyfish网络不同层级学到的特征价值完全不一样。以在 ImageNet 上预训练的 ConvNeXt 为例底层靠近输入的 stage0、stem学到的是边缘、纹理、颜色等通用视觉特征。这些特征是所有视觉任务的基础工具几乎不需要针对特定业务修改改了反而会破坏通用表达能力中层stage1、stage2学到的是部件、纹理组合等中级特征有一定通用性只需要小幅调整高层stage3、最终归一化层学到的是高级语义特征和预训练任务强相关需要更多地向业务数据对齐分类头head完全是为预训练类别设计的需要全部重新学习更新幅度应该最大。如果给所有层设置相同的学习率就会出现底层改过头、高层改不够的尴尬局面。而 LLRD 的思路就是给不同层级匹配不同的更新幅度让底层慢更、顶层更快。分层学习率衰减Layer-wise Learning Rate Decay, LLRD是从模型输出层到输入层学习率按照固定的衰减系数逐级递减越靠近分类头学习率越高越靠近输入学习率越低。它是一种空间维度的学习率差异化策略——同一训练时刻不同层的学习率不同层间的比例关系保持固定。【LLRD 逐层学习率衰减分组结果】[0]head/weight|LR5.00e-04|WD5e-04|params1,024[1]head/bias|LR5.00e-04|WD0e00|params2[2]norm|LR4.00e-04|WD0e00|params1,024[3]stages[3]/weight|LR4.00e-04|WD5e-04|params2,122,240[4]stages[3]/bias|LR4.00e-04|WD0e00|params4,096[5]downsample[3]/weight|LR3.20e-04|WD5e-04|params524,288[6]downsample[3]/bias|LR3.20e-04|WD0e00|params1,024[7]stages[2]/weight|LR2.56e-04|WD5e-04|params536,832[8]stages[2]/bias|LR2.56e-04|WD0e00|params2,048[9]downsample[2]/weight|LR2.05e-04|WD5e-04|params131,072[10]downsample[2]/bias|LR2.05e-04|WD0e00|params512[11]stages[1]/weight|LR1.64e-04|WD5e-04|params137,344[12]stages[1]/bias|LR1.64e-04|WD0e00|params1,024[13]downsample[1]/weight|LR1.31e-04|WD5e-04|params32,768[14]downsample[1]/bias|LR1.31e-04|WD0e00|params256[15]stages[0]/weight|LR1.05e-04|WD5e-04|params35,904[16]stages[0]/bias|LR1.05e-04|WD0e00|params512[17]downsample[0]/weight|LR8.39e-05|WD5e-04|params3,072[18]downsample[0]/bias|LR8.39e-05|WD0e00|params70----------------------------------------------------------------------importtorchimporttorch.nnasnn# 超参数 BASE_LR5e-4# 顶层基准学习率LLRD_DECAY0.8# 逐层衰减系数WEIGHT_DECAY5e-4# 权重衰减系数NUM_CLASSES2# 分类类别数# 基础组件classConvBlock(nn.Module):ConvNeXt 基础块def__init__(self,dim):super().__init__()self.dwconvnn.Conv2d(dim,dim,7,padding3,groupsdim)self.normnn.LayerNorm(dim)self.pwconv1nn.Linear(dim,4*dim)self.actnn.GELU()self.pwconv2nn.Linear(4*dim,dim)defforward(self,x):inputx xself.dwconv(x)# 通道转到最后一维适配 LayerNormxx.permute(0,2,3,1).contiguous()xself.norm(x)xself.pwconv1(x)xself.act(x)xself.pwconv2(x)# 通道还原回 (B,C,H,W)xx.permute(0,3,1,2).contiguous()returnxinputclassDownsample(nn.Module):下采样层LayerNorm 卷积def__init__(self,in_dim,out_dim,stride2):super().__init__()self.normnn.LayerNorm(in_dim)self.convnn.Conv2d(in_dim,out_dim,kernel_sizestride,stridestride)defforward(self,x):# x 形状: (B, C, H, W) → 转成 (B, H, W, C) 过 LayerNormxx.permute(0,2,3,1).contiguous()xself.norm(x)# 还原回卷积需要的 (B, C, H, W)xx.permute(0,3,1,2).contiguous()xself.conv(x)returnx# 模拟一个神经网络classSimpleModel(nn.Module):def__init__(self,num_classes2):super().__init__()# 4个下采样层 4个stage与你的业务模型结构一一对应self.downsample_layersnn.ModuleList([Downsample(3,64,stride4),# 第0层stem输入3通道Downsample(64,128,stride2),# 第1层stage0→stage1 下采样Downsample(128,256,stride2),# 第2层stage1→stage2 下采样Downsample(256,512,stride2),# 第3层stage2→stage3 下采样])self.stagesnn.ModuleList([nn.Sequential(ConvBlock(64)),nn.Sequential(ConvBlock(128)),nn.Sequential(ConvBlock(256)),nn.Sequential(ConvBlock(512)),])self.normnn.LayerNorm(512)# 最终归一化层self.headnn.Linear(512,num_classes)# 分类头defforward(self,x):# 顺序下采样 → 特征提取共4组foriinrange(4):xself.downsample_layers[i](x)xself.stages[i](x)# 全局平均池化xx.mean(dim[2,3])xself.norm(x)xself.head(x)returnx# LLRD 逐层学习率衰减 defbuild_llrd_param_groups(model): 功能为 ConvNeXt 构建逐层学习率衰减参数分组 规则从顶层分类头到底层输入学习率按 decay 系数逐级递减 权重参数加 weight_decaybias/LayerNorm 不加 weight_decay param_groups[]# 最终传给优化器的分组base_lrs[]# 每组基础学习率group_names[]# 分组名称仅打印用defadd_group(params,lr,wd,name):param_groups.append({params:params,lr:lr,weight_decay:wd})base_lrs.append(lr)group_names.append(name)# 初始学习率 基准学习率顶层最大lrBASE_LR# ---------- 1. 分类头 head最高学习率----------head_w,head_b[],[]forname,paraminmodel.head.named_parameters():ifbiasinnameorisinstance(param,nn.LayerNorm):head_b.append(param)else:head_w.append(param)ifhead_w:add_group(head_w,lr,WEIGHT_DECAY,head/weight)ifhead_b:add_group(head_b,lr,0.0,head/bias)# ---------- 2. 最终归一化层 norm ----------lr*LLRD_DECAY add_group(list(model.norm.parameters()),lr,0.0,norm)# ---------- 3. 从深到浅遍历 stages downsample ----------foriinreversed(range(4)):# 处理当前 stagestage_w,stage_b[],[]forname,paraminmodel.stages[i].named_parameters():ifbiasinnameornorminname:stage_b.append(param)else:stage_w.append(param)ifstage_w:add_group(stage_w,lr,WEIGHT_DECAY,fstages[{i}]/weight)ifstage_b:add_group(stage_b,lr,0.0,fstages[{i}]/bias)# 第一次衰减stage → downsamplelr*LLRD_DECAY# 处理当前下采样层ds_w,ds_b[],[]forname,paraminmodel.downsample_layers[i].named_parameters():ifbiasinnameornorminname:ds_b.append(param)else:ds_w.append(param)ifds_w:add_group(ds_w,lr,WEIGHT_DECAY,fdownsample[{i}]/weight)ifds_b:add_group(ds_b,lr,0.0,fdownsample[{i}]/bias)# 第二次衰减进入下一个更浅的 stagelr*LLRD_DECAY# ---------- 打印分组结果 ----------print(*70)print(【LLRD 逐层学习率衰减分组结果】)print(*70)total_params0fori,(pg,lr_val,name)inenumerate(zip(param_groups,base_lrs,group_names)):param_countsum(p.numel()forpinpg[params])total_paramsparam_countprint(f[{i:2d}]{name:20s}| LR{lr_val:.2e}| WD{pg[weight_decay]:.0e}| params{param_count:,})print(-*70)print(f总参数量:{total_params:,})print(f顶层/底层学习率比值:{base_lrs[0]/base_lrs[-1]:.2f})print(*70)returnparam_groups,base_lrs# 运行演示 if__name____main__:# 1. 初始化模型modelSimpleModel(num_classesNUM_CLASSES)# 2. 调用 LLRD 生成分层参数组param_groups,base_lrsbuild_llrd_param_groups(model)# 3. 传入优化器optimizertorch.optim.AdamW(param_groups)# 4. 跑一步训练dummy_inputtorch.randn(2,3,64,64)# batch2, 3通道, 64x64dummy_labeltorch.randint(0,NUM_CLASSES,(2,))loss_fnnn.CrossEntropyLoss()optimizer.zero_grad()outmodel(dummy_input)lossloss_fn(out,dummy_label)loss.backward()optimizer.step()print(f\nLLRD 配置完成训练步正常运行)print(f 输出形状:{out.shape}| 损失值:{loss.item():.4f})