PyTorch 2.x 参数分组优化实战精细化控制学习率与权重衰减深度学习中模型的不同部分往往需要不同的训练策略。PyTorch的参数分组功能Per-parameter options让我们能够为模型的不同组件设置独立的学习率、权重衰减等超参数。这项技术在微调预训练模型、多任务学习等场景中尤为重要。1. 参数分组基础与核心原理参数分组的本质是将模型的参数划分为不同的组每个组可以拥有独立的优化配置。PyTorch的优化器通过param_groups属性管理这些分组每个分组是一个字典包含以下关键信息params: 该组的参数列表lr: 该组的学习率weight_decay: 该组的权重衰减系数其他优化器特定参数如momentum等import torch import torch.nn as nn import torch.optim as optim # 定义一个简单的模型 class CustomModel(nn.Module): def __init__(self): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 64, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2) ) self.classifier nn.Linear(64*16*16, 10) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) x self.classifier(x) return x参数分组的核心优势在于精细化控制不同层/模块可以使用不同的学习策略训练稳定性预训练部分可以使用更小的学习率资源效率冻结部分参数可以减少计算量2. 三种典型场景的参数分组配置2.1 迁移学习中的骨干网络微调在迁移学习中我们通常保持预训练骨干网络backbone的低学习率而为新添加的分类头使用较高的学习率。model CustomModel() pretrained_dict torch.load(pretrained.pth) model.load_state_dict(pretrained_dict, strictFalse) # 冻结特征提取层 for param in model.features.parameters(): param.requires_grad False optimizer optim.Adam([ {params: model.features.parameters(), lr: 1e-5}, # 预训练部分小学习率 {params: model.classifier.parameters(), lr: 1e-3} # 新分类头大学习率 ]) # 训练过程中可以部分解冻 def unfreeze_layers(model, epoch): if epoch 5: for param in model.features[-2:].parameters(): # 解冻最后两层 param.requires_grad True这种配置的典型表现前几个epoch主要训练分类头后期逐渐微调骨干网络的高层特征避免了破坏预训练获得的底层特征2.2 多任务学习中的差异化配置当模型需要同时处理多个相关任务时不同任务对应的模块可能需要不同的优化策略。class MultiTaskModel(nn.Module): def __init__(self): super().__init__() self.shared_backbone nn.Sequential(...) self.task1_head nn.Linear(256, 10) self.task2_head nn.Linear(256, 5) optimizer optim.SGD([ {params: model.shared_backbone.parameters()}, {params: model.task1_head.parameters(), lr: 0.01, weight_decay: 1e-4}, {params: model.task2_head.parameters(), lr: 0.005, momentum: 0.8} ], lr0.001, momentum0.9)关键考虑因素任务重要性不同可设置不同学习率简单任务可以使用更强的正则化复杂任务可能需要更稳定的优化策略2.3 特殊层的独立参数配置某些特殊层如BatchNorm通常需要与其他层区别对待def get_optimizer(model): params [] for name, param in model.named_parameters(): if bn in name or bias in name: # BatchNorm和bias通常使用更大学习率和无weight decay params.append({params: param, weight_decay: 0, lr_mult: 2}) else: params.append({params: param}) base_lr 0.001 optimizer optim.AdamW([ {**group, lr: base_lr * group.get(lr_mult, 1)} for group in params ], weight_decay1e-4) return optimizer这种配置的科学依据BatchNorm的scale/shift参数本质上是线性变换bias项不需要L2正则化这些参数通常需要更大的更新幅度3. 参数分组的高级技巧与调试3.1 动态参数调整策略训练过程中动态调整不同组的参数可以进一步提升效果def adjust_group_lr(optimizer, epoch): 随着训练进程调整不同组的学习率 for i, group in enumerate(optimizer.param_groups): if i 0: # 骨干网络 group[lr] 0.01 * (0.1 ** (epoch // 10)) else: # 分类头 group[lr] 0.1 * (0.95 ** epoch)3.2 梯度裁剪的组别化应用不同参数组可能需要不同的梯度裁剪阈值def group_clip_grad(optimizer, max_norm): for group in optimizer.param_groups: nn.utils.clip_grad_norm_( group[params], max_norm * group.get(clip_mult, 1) )3.3 参数组的状态监控调试时监控各组的状态很有帮助def print_group_stats(optimizer): for i, group in enumerate(optimizer.param_groups): grads [p.grad.norm().item() for p in group[params] if p.grad is not None] print(fGroup {i}: lr{group[lr]}, avg_grad{sum(grads)/len(grads):.4f})4. 完整实战案例CIFAR-10上的ResNet-18微调下面是一个完整的参数分组实战示例import torch import torchvision from torch import nn, optim from torch.utils.data import DataLoader from torchvision.transforms import Compose, ToTensor, Normalize # 1. 准备数据 transform Compose([ ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) trainset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) trainloader DataLoader(trainset, batch_size128, shuffleTrue) # 2. 加载预训练模型并修改 model torchvision.models.resnet18(pretrainedTrue) model.fc nn.Linear(model.fc.in_features, 10) # 替换最后的全连接层 # 3. 设置参数分组 optimizer optim.SGD([ {params: [p for n, p in model.named_parameters() if fc not in n], lr: 0.001}, {params: model.fc.parameters(), lr: 0.01} ], momentum0.9, weight_decay5e-4) # 4. 训练循环 criterion nn.CrossEntropyLoss() for epoch in range(20): for inputs, labels in trainloader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() # 对骨干网络使用较小的梯度更新 with torch.no_grad(): for name, param in model.named_parameters(): if fc not in name: param.grad * 0.1 optimizer.step()这个案例展示了骨干网络使用小学习率(0.001)和小梯度更新(×0.1)新分类头使用大学习率(0.01)和完整梯度所有参数共享momentum和weight_decay5. 参数分组的可视化分析理解参数分组的效果可以通过多种可视化方式学习率变化曲线import matplotlib.pyplot as plt lrs {fgroup_{i}: [] for i in range(len(optimizer.param_groups))} for epoch in range(20): adjust_group_lr(optimizer, epoch) for i, group in enumerate(optimizer.param_groups): lrs[fgroup_{i}].append(group[lr]) plt.figure(figsize(10, 5)) for name, values in lrs.items(): plt.plot(values, labelname) plt.legend() plt.title(Learning Rate Schedule by Group) plt.xlabel(Epoch) plt.ylabel(Learning Rate)梯度分布对比def plot_grad_distribution(model): gradients {name: param.grad.abs().mean().item() for name, param in model.named_parameters() if param.grad is not None} plt.figure(figsize(12, 6)) plt.bar(range(len(gradients)), list(gradients.values())) plt.xticks(range(len(gradients)), list(gradients.keys()), rotation90) plt.title(Average Gradient Magnitude by Layer) plt.ylabel(Gradient Magnitude)这些可视化可以帮助我们确认不同组确实按照预期接收不同的更新强度发现潜在的问题如某些层梯度消失/爆炸优化分组策略和超参数设置