15 PyTorch 小白从零学教程系列:只生成模型训练之训练循环的实现
在前一篇文章中,我们讨论了模型训练之选择优化器的内容,了解到不同的优化器在训练过程中对模型参数更新的方式有所不同。在这一篇中,我们将深入探讨如何实现一个完整的训练循环,以便在选定优化器的基础上进行模型训练。
训练循环的基本概念
训练循环是深度学习中最核心的部分之一,它负责执行模型的前向传播、计算损失、后向传播以及更新模型参数的过程。具体而言,一个训练循环一般包含以下几个步骤:
- 加载数据:从数据集加载一个批次的数据。
- 前向传播:将输入数据传递给模型,得到预测结果。
- 计算损失:与真实标签进行比较,计算损失值。
- 反向传播:计算损失的梯度。
- 优化器更新参数:根据计算得到的梯度更新模型的参数。
- 记录和输出:记录损失和性能指标,输出训练信息。
示例:实现训练循环
接下来,我们将通过具体的案例来实现一个简单的训练循环。假设我们正在训练一个用于图像分类的模型,使用的是 CIFAR-10
数据集。
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()
则执行多步优化更新当前参数。
结论
本篇文章通过一个具体的案例,详细介绍了如何实现模型训练中的训练循环。将前一篇的优化器选择与当前的训练循环结合,我们能够全面理解模型训练的关键步骤。接下来,我们将在下一篇文章中深入讨论如何评估模型性能与调优。
保持对训练过程的关注,不断进行改进与提升,这是深度学习中的一项重要技能。继续学习,迎接下一篇文章的挑战吧!