15 PyTorch 小白从零学教程系列：只生成模型训练之训练循环的实现

查看大图

训练循环看似重复，但顺序很重要：前向计算、计算 loss、清梯度、反向传播、更新参数，再记录指标。

官方教程：PyTorch optimization

查看大图

我会在每个 epoch 记录训练 loss、验证 loss 和学习率。没有日志，训练失败时只能猜。

在前一篇文章中，我们讨论了模型训练之选择优化器的内容，了解到不同的优化器在训练过程中对模型参数更新的方式有所不同。在这一篇中，我们将深入探讨如何实现一个完整的训练循环，以便在选定优化器的基础上进行模型训练。

训练循环的基本概念

训练循环是深度学习中最核心的部分之一，它负责执行模型的前向传播、计算损失、后向传播以及更新模型参数的过程。具体而言，一个训练循环一般包含以下几个步骤：

查看大图

实现 PyTorch 训练循环时，先看 zero_grad、forward、loss、backward、step、指标记录和验证模式。

1. 加载数据：从数据集加载一个批次的数据。
2. 前向传播：将输入数据传递给模型，得到预测结果。
3. 计算损失：与真实标签进行比较，计算损失值。
4. 反向传播：计算损失的梯度。
5. 优化器更新参数：根据计算得到的梯度更新模型的参数。
6. 记录和输出：记录损失和性能指标，输出训练信息。

示例：实现训练循环

接下来，我们将通过具体的案例来实现一个简单的训练循环。假设我们正在训练一个用于图像分类的模型，使用的是 CIFAR-10 数据集。

查看大图

学习《PyTorch 小白从零学教程系列：只生成模型训练之训练循环的实现》不必一口气吃完所有细节。先挑一个能动手验证的小问题，再顺着图和正文补齐概念。

1. 准备工作

首先，我们需要安装 PyTorch 和 torchvision，并导入必要的库：

pip install torch torchvision

然后，我们可以开始编写我们的代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 设置随机种子以确保结果可复现
torch.manual_seed(0)

# 超参数
learning_rate = 0.001
num_epochs = 5
batch_size = 64

# 数据预处理
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.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 = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)

2. 定义模型

这里我们定义一个简单的卷积神经网络模型：

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)  # 展平
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 实例化模型
model = SimpleCNN()

3. 训练循环实现

下面是训练循环的具体实现：

# 选择优化器
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 损失函数
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    running_loss = 0.0
  
    for i, (inputs, labels) in enumerate(trainloader):
        # 将输入和标签移动到GPU (如果有的话)
        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()
        if i % 100 == 99:  # 每100个小批次输出一次
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(trainloader)}], Loss: {running_loss / 100:.4f}')
            running_loss = 0.0

在这个循环中，optimizer.zero_grad() 是用来清除前一个小批次的梯度，loss.backward() 将计算得到的梯度赋值给相应的参数，optimizer.step() 则执行多步优化更新当前参数。

查看大图

如果《PyTorch 小白从零学教程系列：只生成模型训练之训练循环的实现》还没完全消化，可以从这张卡片的四个动作重新走一遍。

查看大图

回看《PyTorch 小白从零学教程系列：只生成模型训练之训练循环的实现》时，不必一次做大项目，先用一条简单样例确认主线是否清楚。

结论

本篇文章通过一个具体的案例，详细介绍了如何实现模型训练中的训练循环。将前一篇的优化器选择与当前的训练循环结合，我们能够全面理解模型训练的关键步骤。接下来，我们将在下一篇文章中深入讨论如何评估模型性能与调优。

保持对训练过程的关注，不断进行改进与提升，这是深度学习中的一项重要技能。继续学习，迎接下一篇文章的挑战吧！