PyTorch学习率调度实战：用CosineAnnealingLR和WarmRestarts搞定图像分类任务（以ResNet18为例）

PyTorch学习率调度实战用CosineAnnealingLR和WarmRestarts搞定图像分类任务以ResNet18为例在深度学习模型训练中学习率调度策略的选择往往决定了模型能否快速收敛到最优解。想象一下你正在训练一个ResNet18模型进行图像分类前几轮训练进展顺利但到了中后期验证集准确率却停滞不前——这可能正是学习率调度策略需要优化的信号。本文将带你深入实战通过PyTorch的CosineAnnealingLR和CosineAnnealingWarmRestarts两种调度器为ResNet18模型打造一个高效的学习率调节方案。1. 环境准备与数据加载在开始之前我们需要准备好实验环境。这里推荐使用Python 3.8和PyTorch 1.10版本确保能够支持最新的调度器功能。首先安装必要的依赖pip install torch torchvision matplotlib tensorboard对于图像分类任务我们选择CIFAR-10数据集作为示例。这个数据集包含10个类别的6万张32x32彩色图像非常适合验证ResNet18这样的轻量级模型import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 数据增强和归一化 train_transform transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, padding4), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) # 加载数据集 train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) test_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtest_transform) # 创建数据加载器 batch_size 128 train_loader DataLoader(train_set, batch_sizebatch_size, shuffleTrue, num_workers4) test_loader DataLoader(test_set, batch_sizebatch_size, shuffleFalse, num_workers4)2. 模型构建与基础训练流程我们使用PyTorch内置的ResNet18模型并针对CIFAR-10的32x32输入尺寸进行适当调整import torch.nn as nn from torchvision.models import resnet18 def create_model(num_classes10): model resnet18(pretrainedFalse) # 调整第一层卷积适应CIFAR-10的32x32输入 model.conv1 nn.Conv2d(3, 64, kernel_size3, stride1, padding1, biasFalse) # 移除最后的全连接层替换为适合类别数的层 model.fc nn.Linear(model.fc.in_features, num_classes) return model model create_model().cuda()基础训练流程包含损失函数和优化器的设置。这里我们使用交叉熵损失和SGD优化器criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.1, momentum0.9, weight_decay5e-4)3. CosineAnnealingLR调度器实战CosineAnnealingLR实现了一个简单的余弦退火策略学习率在给定的周期内按照余弦曲线从初始值下降到最小值。3.1 参数配置与初始化关键参数说明T_max: 半周期长度epoch数eta_min: 学习率最小值对于CIFAR-10训练通常设置150-200个epoch。我们选择T_max50意味着每50个epoch完成半个余弦周期from torch.optim.lr_scheduler import CosineAnnealingLR scheduler_cosine CosineAnnealingLR(optimizer, T_max50, eta_min1e-5)3.2 训练循环集成将调度器集成到训练循环中每个epoch后调用scheduler.step()def train_with_scheduler(model, train_loader, test_loader, scheduler, epochs150): train_losses [] test_accuracies [] learning_rates [] for epoch in range(epochs): model.train() running_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.cuda(), labels.cuda() optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 记录当前学习率 current_lr optimizer.param_groups[0][lr] learning_rates.append(current_lr) # 更新学习率 scheduler.step() # 评估模型 test_acc evaluate(model, test_loader) test_accuracies.append(test_acc) train_loss running_loss / len(train_loader) train_losses.append(train_loss) print(fEpoch {epoch1}/{epochs} - Loss: {train_loss:.4f}, Acc: {test_acc:.4f}, LR: {current_lr:.6f}) return train_losses, test_accuracies, learning_rates3.3 学习率变化可视化训练完成后我们可以绘制学习率变化曲线import matplotlib.pyplot as plt def plot_learning_rates(learning_rates, title): plt.figure(figsize(10, 5)) plt.plot(learning_rates) plt.xlabel(Epoch) plt.ylabel(Learning Rate) plt.title(title) plt.grid(True) plt.show() # 假设已经完成了训练 # plot_learning_rates(cosine_lrs, CosineAnnealingLR Learning Rate Schedule)典型的余弦退火学习率曲线会呈现周期性波动在50个epoch内从初始值平滑下降到最小值然后重新开始。4. CosineAnnealingWarmRestarts调度器实战CosineAnnealingWarmRestarts是CosineAnnealingLR的改进版在每次重启时保留部分之前的学习率变化记忆通常能带来更好的性能。4.1 参数配置与初始化关键参数说明T_0: 第一次重启的epoch数T_mult: 重启周期倍增因子eta_min: 学习率最小值我们设置初始周期T_030倍增因子T_mult2意味着第一个周期30个epoch第二个周期60个epoch第三个周期120个epoch以此类推...from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts scheduler_warm CosineAnnealingWarmRestarts(optimizer, T_030, T_mult2, eta_min1e-5)4.2 训练循环集成训练循环与之前类似只是更换了调度器# 使用相同的train_with_scheduler函数传入不同的scheduler train_losses_warm, test_accuracies_warm, learning_rates_warm train_with_scheduler( model, train_loader, test_loader, scheduler_warm, epochs150 )4.3 学习率变化对比将两种调度器的学习率曲线绘制在一起对比def compare_schedulers(lr_cosine, lr_warm, epochs): plt.figure(figsize(12, 6)) plt.plot(range(epochs), lr_cosine, labelCosineAnnealingLR (T_max50)) plt.plot(range(epochs), lr_warm, labelCosineAnnealingWarmRestarts (T_030, T_mult2)) plt.xlabel(Epoch) plt.ylabel(Learning Rate) plt.title(Comparison of Learning Rate Schedules) plt.legend() plt.grid(True) plt.show() # 假设已经获取了两种调度器的学习率记录 # compare_schedulers(cosine_lrs, warm_lrs, epochs150)5. 性能分析与实践建议在实际项目中我们需要综合考虑训练时间、资源消耗和模型性能。以下是两种调度器的对比分析特性CosineAnnealingLRCosineAnnealingWarmRestarts收敛速度中等较快最终准确率良好优秀超参数敏感性中等较低计算开销低略高适合场景小中型数据集大型复杂数据集实践中的几点建议初始学习率选择对于ResNet18和CIFAR-100.1是一个不错的起点T_max/T_0设置通常设为总epoch数的1/3到1/2eta_min设置一般设为初始学习率的1/100到1/1000Warm Restarts优势在训练后期能帮助模型跳出局部最优# 一个综合了最佳实践的配置示例 optimizer torch.optim.SGD(model.parameters(), lr0.1, momentum0.9, weight_decay5e-4) scheduler CosineAnnealingWarmRestarts(optimizer, T_040, T_mult2, eta_min1e-4)在CIFAR-10上的实验表明使用Warm Restarts的调度器通常能比固定学习率或简单余弦退火获得1-2%的准确率提升。更重要的是这种提升不需要额外的训练时间或计算资源只需要正确配置调度器参数。