PyTorch 2.0 CIFAR-10 模型优化：3种学习率策略对比与实战指南

1. 引言：为什么学习率策略如此关键？

在深度学习模型的训练过程中，学习率（Learning Rate）无疑是最关键的超参数之一。它决定了模型参数在每次迭代中更新的步长大小，直接影响着模型的收敛速度和最终性能。对于CIFAR-10这样的经典图像分类任务，选择合适的学习率策略往往能让模型准确率提升5-10个百分点。

想象一下这样的场景：你精心设计了一个CNN模型，数据预处理也做得非常规范，但训练过程中要么收敛缓慢，要么在测试集上表现平平。这时候，问题很可能就出在学习率策略上。学习率太大可能导致模型在最优解附近震荡甚至发散；太小则会让训练过程变得极其缓慢，甚至陷入局部最优。

PyTorch 2.0为我们提供了多种内置的学习率调度器（Learning Rate Scheduler），本文将重点对比三种最常用的策略：StepLR（阶梯式衰减）、ExponentialLR（指数衰减）和CosineAnnealingLR（余弦退火）。通过完整的代码实现和量化对比，帮助你在CIFAR-10分类任务中将测试准确率提升至69.78%甚至更高。

2. 实验环境与基准模型配置

2.1 实验环境准备

在开始对比实验前，我们需要确保所有参与者站在同一起跑线上。以下是本次实验的环境配置：

import torch import torchvision import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler from torchvision import datasets, transforms # 检查PyTorch版本和设备 print(f"PyTorch版本: {torch.__version__}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"使用设备: {device}")

硬件建议：

• GPU: NVIDIA GTX 1660 或更高（6GB显存足够）
• 内存: 16GB以上
• PyTorch版本: 2.0+

2.2 基准模型架构

我们采用一个中等复杂度的CNN架构作为基准模型，它包含：

class CIFAR10Model(nn.Module): def __init__(self): super(CIFAR10Model, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(128 * 4 * 4, 512), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(512, 10), ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x

模型关键参数：

• 输入尺寸: 3x32x32 (CIFAR-10标准尺寸)
• 卷积层: 3层，通道数分别为32/64/128
• 全连接层: 2层，隐藏单元512个
• Dropout: 0.5 (防止过拟合)

2.3 数据预处理与加载

统一的数据处理流程对公平对比至关重要：

# 数据预处理 transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, padding=4), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) # 加载数据集 train_dataset = datasets.CIFAR10( root='./data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10( root='./data', train=False, download=True, transform=transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=128, shuffle=True, num_workers=2) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=100, shuffle=False, num_workers=2)

数据增强策略：

• 随机水平翻转 (RandomHorizontalFlip)
• 随机裁剪 (RandomCrop with padding)
• 标准化处理 (Normalize)

3. 三种学习率策略详解与实现

3.1 StepLR：阶梯式学习率衰减

原理： StepLR每隔固定的epoch数将学习率乘以一个衰减系数gamma。这种策略简单直观，适合训练初期快速下降，后期精细调整的场景。

def train_with_steplr(epochs=50, lr=0.1, gamma=0.1, step_size=20): model = CIFAR10Model().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma) train_losses, test_accs = [], [] for epoch in range(epochs): # 训练阶段 model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 更新学习率 scheduler.step() # 测试阶段 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() # 记录结果 train_loss = running_loss / len(train_loader) test_acc = 100 * correct / total train_losses.append(train_loss) test_accs.append(test_acc) print(f'Epoch {epoch+1}/{epochs} | LR: {scheduler.get_last_lr()[0]:.6f} | ' f'Train Loss: {train_loss:.4f} | Test Acc: {test_acc:.2f}%') return test_accs[-1] # 返回最终测试准确率

参数选择建议：

• 初始学习率 (lr): 0.1 (SGD常用)
• 衰减系数 (gamma): 0.1 (每次衰减为原来的1/10)
• 衰减步长 (step_size): 20 (每20个epoch衰减一次)

3.2 ExponentialLR：指数衰减学习率

原理： ExponentialLR每个epoch都按指数规律衰减学习率，衰减速度由gamma参数控制。相比StepLR，它的衰减更加平滑连续。

def train_with_exponentiallr(epochs=50, lr=0.1, gamma=0.95): model = CIFAR10Model().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=gamma) train_losses, test_accs = [], [] for epoch in range(epochs): # 训练阶段 model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 更新学习率 scheduler.step() # 测试阶段 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() # 记录结果 train_loss = running_loss / len(train_loader) test_acc = 100 * correct / total train_losses.append(train_loss) test_accs.append(test_acc) print(f'Epoch {epoch+1}/{epochs} | LR: {scheduler.get_last_lr()[0]:.6f} | ' f'Train Loss: {train_loss:.4f} | Test Acc: {test_acc:.2f}%') return test_accs[-1]

参数选择建议：

• 初始学习率 (lr): 0.1
• 衰减系数 (gamma): 0.95 (每个epoch学习率乘以0.95)
• 特点：适合希望学习率持续缓慢下降的场景

3.3 CosineAnnealingLR：余弦退火学习率

原理： CosineAnnealingLR按照余弦函数的规律调整学习率，从初始值下降到接近0，然后再重新开始。这种周期性变化有助于模型跳出局部最优。

def train_with_cosinelr(epochs=50, lr=0.1, T_max=10, eta_min=0): model = CIFAR10Model().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=T_max, eta_min=eta_min) train_losses, test_accs = [], [] for epoch in range(epochs): # 训练阶段 model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 更新学习率 scheduler.step() # 测试阶段 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() # 记录结果 train_loss = running_loss / len(train_loader) test_acc = 100 * correct / total train_losses.append(train_loss) test_accs.append(test_acc) print(f'Epoch {epoch+1}/{epochs} | LR: {scheduler.get_last_lr()[0]:.6f} | ' f'Train Loss: {train_loss:.4f} | Test Acc: {test_acc:.2f}%') return test_accs[-1]

参数选择建议：

• 初始学习率 (lr): 0.1
• T_max: 10 (余弦周期长度，单位是epoch)
• eta_min: 0 (学习率下限)
• 特点：适合希望周期性调整学习率的场景

4. 实验结果对比与分析

4.1 量化结果对比

我们固定随机种子，在相同条件下运行三种策略各50个epoch，得到如下结果：

关键发现：

1. ExponentialLR在本实验中表现最佳，达到69.78%的测试准确率
2. CosineAnnealingLR虽然最终准确率略低，但训练过程更加稳定
3. StepLR在第一次学习率下降后性能提升明显

4.2 学习率变化曲线对比

import matplotlib.pyplot as plt # 模拟三种策略的学习率变化 epochs = 50 step_lr = [0.1 * (0.1 ** (i // 20)) for i in range(epochs)] exp_lr = [0.1 * (0.95 ** i) for i in range(epochs)] cos_lr = [0.05 + 0.05 * (1 + math.cos(math.pi * i / 10)) for i in range(epochs)] plt.figure(figsize=(10, 6)) plt.plot(step_lr, label='StepLR (gamma=0.1, step_size=20)') plt.plot(exp_lr, label='ExponentialLR (gamma=0.95)') plt.plot(cos_lr, label='CosineAnnealingLR (T_max=10)') plt.xlabel('Epoch') plt.ylabel('Learning Rate') plt.title('Learning Rate Schedule Comparison') plt.legend() plt.grid(True) plt.show()

![学习率变化曲线对比图]

4.3 训练动态分析

StepLR：

• 优点：简单直观，计算开销小
• 缺点：学习率突变可能导致训练不稳定
• 适用场景：资源有限，需要简单有效策略的情况

ExponentialLR：

• 优点：平滑衰减，训练稳定
• 缺点：后期学习率可能过小，收敛缓慢
• 适用场景：希望学习率持续缓慢下降的任务

CosineAnnealingLR：

• 优点：周期性变化有助于跳出局部最优
• 缺点：需要更多epoch才能体现优势
• 适用场景：计算资源充足，追求更高模型性能

5. 进阶技巧与实战建议

5.1 学习率预热（Warmup）

对于深层网络，初始阶段可以采用线性warmup策略：

def warmup_scheduler(optimizer, warmup_epochs, initial_lr, target_lr): def lr_lambda(epoch): if epoch < warmup_epochs: return initial_lr + (target_lr - initial_lr) * epoch / warmup_epochs else: return target_lr return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

5.2 组合策略：CosineAnnealing with Restarts

PyTorch提供的CosineAnnealingWarmRestarts结合了余弦退火和周期性重启的优点：

scheduler = lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=10, T_mult=2, eta_min=0)

5.3 自动化学习率调整

PyTorch的ReduceLROnPlateau可以根据验证集表现自动调整学习率：

scheduler = lr_scheduler.ReduceLROnPlateau( optimizer, mode='max', factor=0.1, patience=5)

5.4 实际项目中的经验法则

1. 初始学习率选择：

   • SGD: 0.1-0.01
   • Adam: 0.001-0.0001

2. 衰减策略选择：

   • 小型数据集/简单任务：StepLR或ExponentialLR
   • 大型数据集/复杂任务：CosineAnnealing或组合策略

3. 监控与调整：

   • 始终监控训练/验证损失曲线
   • 如果验证损失波动大，尝试减小学习率或增加batch size
   • 如果验证损失下降缓慢，可以适当增大学习率

6. 完整代码实现与复现指南

6.1 完整训练脚本

import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler from torchvision import datasets, transforms import math import matplotlib.pyplot as plt # 1. 定义模型 class CIFAR10Model(nn.Module): def __init__(self): super(CIFAR10Model, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(128 * 4 * 4, 512), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(512, 10), ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x # 2. 数据加载 transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, padding=4), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) train_dataset = datasets.CIFAR10( root='./data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10( root='./data', train=False, download=True, transform=transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=128, shuffle=True, num_workers=2) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=100, shuffle=False, num_workers=2) # 3. 训练函数 def train_model(scheduler_name, epochs=50): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = CIFAR10Model().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) # 选择学习率策略 if scheduler_name == 'StepLR': scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1) elif scheduler_name == 'ExponentialLR': scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.95) elif scheduler_name == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=10) else: raise ValueError("Unknown scheduler name") history = {'train_loss': [], 'test_acc': [], 'lr': []} for epoch in range(epochs): # 训练阶段 model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 更新学习率 scheduler.step() # 测试阶段 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() # 记录结果 train_loss = running_loss / len(train_loader) test_acc = 100 * correct / total history['train_loss'].append(train_loss) history['test_acc'].append(test_acc) history['lr'].append(scheduler.get_last_lr()[0]) print(f'{scheduler_name} Epoch {epoch+1}/{epochs} | ' f'LR: {history["lr"][-1]:.6f} | ' f'Train Loss: {train_loss:.4f} | Test Acc: {test_acc:.2f}%') return history # 4. 运行实验 step_history = train_model('StepLR') exp_history = train_model('ExponentialLR') cos_history = train_model('CosineAnnealingLR') # 5. 可视化结果 plt.figure(figsize=(12, 8)) plt.subplot(2, 2, 1) plt.plot(step_history['test_acc'], label='StepLR') plt.plot(exp_history['test_acc'], label='ExponentialLR') plt.plot(cos_history['test_acc'], label='CosineAnnealingLR') plt.xlabel('Epoch') plt.ylabel('Test Accuracy (%)') plt.title('Test Accuracy Comparison') plt.legend() plt.grid(True) plt.subplot(2, 2, 2) plt.plot(step_history['lr'], label='StepLR') plt.plot(exp_history['lr'], label='ExponentialLR') plt.plot(cos_history['lr'], label='CosineAnnealingLR') plt.xlabel('Epoch') plt.ylabel('Learning Rate') plt.title('Learning Rate Schedule') plt.legend() plt.grid(True) plt.subplot(2, 2, 3) plt.plot(step_history['train_loss'], label='StepLR') plt.plot(exp_history['train_loss'], label='ExponentialLR') plt.plot(cos_history['train_loss'], label='CosineAnnealingLR') plt.xlabel('Epoch') plt.ylabel('Training Loss') plt.title('Training Loss Comparison') plt.legend() plt.grid(True) plt.tight_layout() plt.show()

6.2 复现注意事项

1. 随机种子固定：

   torch.manual_seed(42) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False

2. 硬件要求：

   • GPU显存: ≥4GB (batch_size=128)
   • 训练时间: 约30分钟/50epochs (RTX 3060)

3. 参数调整：

   • 如果显存不足，可以减小batch_size，但需相应调整学习率
   • 想要更高准确率，可以增加epochs到100-200

4. 结果保存：

   torch.save(model.state_dict(), 'cifar10_model.pth')

7. 总结与延伸思考

通过本实验的系统对比，我们验证了不同学习率策略对CIFAR-10分类任务的影响。ExponentialLR在本实验中表现最佳，但这并不意味着它在所有情况下都是最优选择。实际项目中，还需要考虑以下因素：

1. 模型复杂度：更复杂的模型可能需要更精细的学习率调整
2. 数据规模：大数据集可能需要更激进的学习率衰减
3. 优化器选择：Adam等自适应优化器对学习率的敏感度较低

未来优化方向：

1. 尝试组合多种学习率策略（如warmup+cosine）
2. 结合模型架构搜索（NAS）自动优化学习率策略
3. 探索基于强化学习的动态学习率调整方法

在实际项目中，我通常会先用CosineAnnealingLR进行快速实验，然后根据训练曲线再决定是否需要切换到其他策略。记住，没有放之四海而皆准的最佳策略，关键是要理解每种方法背后的原理，然后根据具体问题和数据进行有针对性的选择和调整。