Fine-tuning 学习率策略：分层学习率衰减（Layer-wise Learning Rate Decay, LLRD）
flyfish
网络不同层级学到的特征，价值完全不一样。
以在 ImageNet 上预训练的 ConvNeXt 为例：
底层（靠近输入的 stage0、stem）：学到的是边缘、纹理、颜色等通用视觉特征。这些特征是所有视觉任务的基础工具，几乎不需要针对特定业务修改，改了反而会破坏通用表达能力；
中层（stage1、stage2）：学到的是部件、纹理组合等中级特征，有一定通用性，只需要小幅调整；
高层（stage3、最终归一化层）：学到的是高级语义特征，和预训练任务强相关，需要更多地向业务数据对齐；
分类头（head）：完全是为预训练类别设计的，需要全部重新学习，更新幅度应该最大。
如果给所有层设置相同的学习率，就会出现底层改过头、高层改不够的尴尬局面。而 LLRD 的思路，就是给不同层级匹配不同的更新幅度，让底层慢更、顶层更快。
分层学习率衰减（Layer-wise Learning Rate Decay, LLRD）是从模型输出层到输入层，学习率按照固定的衰减系数逐级递减；越靠近分类头，学习率越高；越靠近输入，学习率越低。
它是一种空间维度的学习率差异化策略——同一训练时刻，不同层的学习率不同，层间的比例关系保持固定。
【LLRD 逐层学习率衰减分组结果】======================================================================[0]head/weight|LR=5.00e-04|WD=5e-04|params=1,024[1]head/bias|LR=5.00e-04|WD=0e+00|params=2[2]norm|LR=4.00e-04|WD=0e+00|params=1,024[3]stages[3]/weight|LR=4.00e-04|WD=5e-04|params=2,122,240[4]stages[3]/bias|LR=4.00e-04|WD=0e+00|params=4,096[5]downsample[3]/weight|LR=3.20e-04|WD=5e-04|params=524,288[6]downsample[3]/bias|LR=3.20e-04|WD=0e+00|params=1,024[7]stages[2]/weight|LR=2.56e-04|WD=5e-04|params=536,832[8]stages[2]/bias|LR=2.56e-04|WD=0e+00|params=2,048[9]downsample[2]/weight|LR=2.05e-04|WD=5e-04|params=131,072[10]downsample[2]/bias|LR=2.05e-04|WD=0e+00|params=512[11]stages[1]/weight|LR=1.64e-04|WD=5e-04|params=137,344[12]stages[1]/bias|LR=1.64e-04|WD=0e+00|params=1,024[13]downsample[1]/weight|LR=1.31e-04|WD=5e-04|params=32,768[14]downsample[1]/bias|LR=1.31e-04|WD=0e+00|params=256[15]stages[0]/weight|LR=1.05e-04|WD=5e-04|params=35,904[16]stages[0]/bias|LR=1.05e-04|WD=0e+00|params=512[17]downsample[0]/weight|LR=8.39e-05|WD=5e-04|params=3,072[18]downsample[0]/bias|LR=8.39e-05|WD=0e+00|params=70----------------------------------------------------------------------importtorchimporttorch.nnasnn# ===================== 超参数 =====================BASE_LR=5e-4# 顶层基准学习率LLRD_DECAY=0.8# 逐层衰减系数WEIGHT_DECAY=5e-4# 权重衰减系数NUM_CLASSES=2# 分类类别数# ===================== 基础组件=====================classConvBlock(nn.Module):"""ConvNeXt 基础块"""def__init__(self,dim):super().__init__()self.dwconv=nn.Conv2d(dim,dim,7,padding=3,groups=dim)self.norm=nn.LayerNorm(dim)self.pwconv1=nn.Linear(dim,4*dim)self.act=nn.GELU()self.pwconv2=nn.Linear(4*dim,dim)defforward(self,x):input=x x=self.dwconv(x)# 通道转到最后一维，适配 LayerNormx=x.permute(0,2,3,1).contiguous()x=self.norm(x)x=self.pwconv1(x)x=self.act(x)x=self.pwconv2(x)# 通道还原回 (B,C,H,W)x=x.permute(0,3,1,2).contiguous()returnx+inputclassDownsample(nn.Module):"""下采样层：LayerNorm + 卷积"""def__init__(self,in_dim,out_dim,stride=2):super().__init__()self.norm=nn.LayerNorm(in_dim)self.conv=nn.Conv2d(in_dim,out_dim,kernel_size=stride,stride=stride)defforward(self,x):# x 形状: (B, C, H, W) → 转成 (B, H, W, C) 过 LayerNormx=x.permute(0,2,3,1).contiguous()x=self.norm(x)# 还原回卷积需要的 (B, C, H, W)x=x.permute(0,3,1,2).contiguous()x=self.conv(x)returnx# ===================== 模拟一个神经网络=====================classSimpleModel(nn.Module):def__init__(self,num_classes=2):super().__init__()# 4个下采样层 + 4个stage，与你的业务模型结构一一对应self.downsample_layers=nn.ModuleList([Downsample(3,64,stride=4),# 第0层：stem，输入3通道Downsample(64,128,stride=2),# 第1层：stage0→stage1 下采样Downsample(128,256,stride=2),# 第2层：stage1→stage2 下采样Downsample(256,512,stride=2),# 第3层：stage2→stage3 下采样])self.stages=nn.ModuleList([nn.Sequential(ConvBlock(64)),nn.Sequential(ConvBlock(128)),nn.Sequential(ConvBlock(256)),nn.Sequential(ConvBlock(512)),])self.norm=nn.LayerNorm(512)# 最终归一化层self.head=nn.Linear(512,num_classes)# 分类头defforward(self,x):# 顺序：下采样 → 特征提取，共4组foriinrange(4):x=self.downsample_layers[i](x)x=self.stages[i](x)# 全局平均池化x=x.mean(dim=[2,3])x=self.norm(x)x=self.head(x)returnx# ===================== LLRD 逐层学习率衰减 =====================defbuild_llrd_param_groups(model):""" 功能：为 ConvNeXt 构建逐层学习率衰减参数分组 规则：从顶层分类头到底层输入，学习率按 decay 系数逐级递减 权重参数加 weight_decay，bias/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)# 初始学习率 = 基准学习率（顶层最大）lr=BASE_LR# ---------- 1. 分类头 head（最高学习率）----------head_w,head_b=[],[]forname,paraminmodel.head.named_parameters():if'bias'innameorisinstance(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():if'bias'innameor'norm'inname:stage_b.append(param)else:stage_w.append(param)ifstage_w:add_group(stage_w,lr,WEIGHT_DECAY,f"stages[{i}]/weight")ifstage_b:add_group(stage_b,lr,0.0,f"stages[{i}]/bias")# 第一次衰减：stage → downsamplelr*=LLRD_DECAY# 处理当前下采样层ds_w,ds_b=[],[]forname,paraminmodel.downsample_layers[i].named_parameters():if'bias'innameor'norm'inname:ds_b.append(param)else:ds_w.append(param)ifds_w:add_group(ds_w,lr,WEIGHT_DECAY,f"downsample[{i}]/weight")ifds_b:add_group(ds_b,lr,0.0,f"downsample[{i}]/bias")# 第二次衰减：进入下一个更浅的 stagelr*=LLRD_DECAY# ---------- 打印分组结果 ----------print("="*70)print("【LLRD 逐层学习率衰减分组结果】")print("="*70)total_params=0fori,(pg,lr_val,name)inenumerate(zip(param_groups,base_lrs,group_names)):param_count=sum(p.numel()forpinpg['params'])total_params+=param_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. 初始化模型model=SimpleModel(num_classes=NUM_CLASSES)# 2. 调用 LLRD 生成分层参数组param_groups,base_lrs=build_llrd_param_groups(model)# 3. 传入优化器optimizer=torch.optim.AdamW(param_groups)# 4. 跑一步训练dummy_input=torch.randn(2,3,64,64)# batch=2, 3通道, 64x64dummy_label=torch.randint(0,NUM_CLASSES,(2,))loss_fn=nn.CrossEntropyLoss()optimizer.zero_grad()out=model(dummy_input)loss=loss_fn(out,dummy_label)loss.backward()optimizer.step()print(f"\nLLRD 配置完成，训练步正常运行")print(f" 输出形状:{out.shape}| 损失值:{loss.item():.4f}")
