PyTorch 2.x Optimizer 参数分组实战：3种场景下的学习率与权重衰减独立配置-开发者社区

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_size=3, padding=1), 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, strict=False) # 冻结特征提取层 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} ], lr=0.001, momentum=0.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_decay=1e-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(f'Group {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', train=True, download=True, transform=transform) trainloader = DataLoader(trainset, batch_size=128, shuffle=True) # 2. 加载预训练模型并修改 model = torchvision.models.resnet18(pretrained=True) 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} ], momentum=0.9, weight_decay=5e-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_decay

5. 参数分组的可视化分析

理解参数分组的效果可以通过多种可视化方式：

学习率变化曲线：

import matplotlib.pyplot as plt lrs = {f'group_{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[f'group_{i}'].append(group['lr']) plt.figure(figsize=(10, 5)) for name, values in lrs.items(): plt.plot(values, label=name) 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()), rotation=90) plt.title('Average Gradient Magnitude by Layer') plt.ylabel('Gradient Magnitude')

这些可视化可以帮助我们：