余弦退火示例pytorch

余弦退火学习率调度器是一种用于动态调整学习率的方法，能够在训练过程中逐渐降低学习率，从而有助于模型更好地收敛。我们可以通过PyTorch中的torch.optim.lr_scheduler.CosineAnnealingLR来实现这种调度器。

下面是一个包含3层Conv2d模型的示例，演示如何使用余弦退火学习率调度器进行训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchvision import datasets, transforms

# 定义简单的3层卷积神经网络
class SimpleConvNet(nn.Module):
    def __init__(self):
        super(SimpleConvNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.fc = nn.Linear(128*7*7, 10)  # 假设输入图像大小为28x28

    def forward(self, x):
        x = nn.ReLU()(self.conv1(x))
        x = nn.MaxPool2d(2)(x)
        x = nn.ReLU()(self.conv2(x))
        x = nn.MaxPool2d(2)(x)
        x = nn.ReLU()(self.conv3(x))
        x = nn.MaxPool2d(2)(x)
        x = x.view(-1, 128*7*7)  # 展平
        x = self.fc(x)
        return x

# 加载数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 创建模型、损失函数和优化器
model = SimpleConvNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

# 设置余弦退火学习率调度器
scheduler = CosineAnnealingLR(optimizer, T_max=10)  # T_max 是一个周期的长度

# 训练过程
num_epochs = 20
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 每个epoch结束后更新学习率
    scheduler.step()
    
    # 打印学习率和损失
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}, Learning Rate: {scheduler.get_last_lr()[0]:.6f}")

print("Training completed.")

代码说明

1. 模型定义：定义了一个简单的3层卷积神经网络，每一层之后接ReLU激活函数和最大池化层。
2. 数据加载：使用MNIST数据集并进行标准化处理。
3. 创建优化器和调度器：使用随机梯度下降（SGD）作为优化器，初始学习率为0.1，动量为0.9。使用CosineAnnealingLR作为学习率调度器，周期长度（T_max）设置为10。
4. 训练过程：在每个epoch结束后调用scheduler.step()来更新学习率，并打印当前的损失值和学习率。