跳到主要内容

PyTorch 训练循环

在深度学习中，训练循环是模型学习的核心过程。它涉及从数据中提取特征、计算损失、更新模型参数等步骤。本文将详细介绍如何在PyTorch中实现一个完整的训练循环，并通过实际案例帮助你理解其工作原理。

什么是训练循环？​

训练循环是指模型在训练数据上反复迭代的过程。每次迭代通常包括以下步骤：

  1. 数据加载：从数据集中获取一批数据。
  2. 前向传播：将数据输入模型，计算输出。
  3. 损失计算：比较模型输出与真实标签，计算损失。
  4. 反向传播：根据损失计算梯度。
  5. 参数更新：使用优化器更新模型参数。

通过多次迭代，模型逐渐学习到数据中的模式，从而提高预测准确性。

训练循环的基本结构​

以下是一个典型的PyTorch训练循环的基本结构：

python
import torch
import torch.nn as nn
import torch.optim as optim

# 假设我们有一个简单的模型
model = nn.Linear(10, 1)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 假设我们有一些数据
inputs = torch.randn(100, 10)
targets = torch.randn(100, 1)

# 训练循环
for epoch in range(100):
    # 前向传播
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

代码解释​

  1. 模型定义：我们使用 nn.Linear 定义了一个简单的线性模型。
  2. 损失函数和优化器：nn.MSELoss() 是均方误差损失函数，optim.SGD 是随机梯度下降优化器。
  3. 数据加载：我们生成了一些随机数据作为输入和目标。
  4. 训练循环：
    • 前向传播：将输入数据传递给模型，计算输出。
    • 损失计算：计算模型输出与目标之间的损失。
    • 反向传播：调用 loss.backward() 计算梯度。
    • 参数更新：调用 optimizer.step() 更新模型参数。

实际案例：手写数字识别​

让我们通过一个实际案例来进一步理解训练循环。我们将使用MNIST数据集训练一个简单的卷积神经网络（CNN）来识别手写数字。

python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 定义简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(32*28*28, 10)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32*28*28)
        x = self.fc1(x)
        return x

model = SimpleCNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item()}')

代码解释​

  1. 数据加载：我们使用 torchvision 加载MNIST数据集，并应用数据转换。
  2. 模型定义：我们定义了一个简单的CNN模型，包含一个卷积层和一个全连接层。
  3. 损失函数和优化器：使用交叉熵损失函数和Adam优化器。
  4. 训练循环：在每个epoch中，我们遍历数据加载器，执行前向传播、损失计算、反向传播和参数更新。

总结​

通过本文，你学习了如何在PyTorch中实现一个完整的训练循环。我们从基本概念入手，逐步讲解了数据加载、前向传播、损失计算、反向传播和参数更新的过程，并通过实际案例展示了如何将这些步骤应用于手写数字识别任务。

提示

为了加深理解，建议你尝试修改代码中的超参数（如学习率、批量大小等），观察它们对训练过程的影响。

附加资源​

练习​

  1. 修改上述代码，使用不同的优化器（如RMSprop）进行训练，并比较结果。
  2. 尝试增加模型的复杂度（如添加更多卷积层），观察训练效果的变化。
  3. 使用其他数据集（如CIFAR-10）进行训练，并调整模型结构以适应新的数据。

通过实践这些练习，你将更好地掌握PyTorch训练循环的实现和应用。