备注
点击此处下载完整示例代码
简介 || 张量 || Autograd || 模型构建 || TensorBoard 支持 || 训练模型 || 模型理解
使用 PyTorch 进行训练 ¶
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
按照下面的视频或 YouTube 上的视频进行操作。
简介
在过去的视频中，我们讨论并演示了：
使用 torch.nn 模块的神经网络层和函数构建模型
自动化梯度计算的原理，这是基于梯度的模型训练的核心
使用 TensorBoard 可视化训练进度和其他活动
在这个视频中，我们将为您添加一些新工具：
我们将熟悉数据集和数据加载器抽象，以及它们如何简化在训练循环中向模型提供数据的过程
我们将讨论具体的损失函数及其使用时机
我们将探讨 PyTorch 优化器，它们实现了根据损失函数的结果调整模型权重的算法
最后，我们将把这些内容整合起来，看看完整的 PyTorch 训练循环是如何运行的。
数据集和数据加载器 ¶
Dataset 和 DataLoader 类封装了从存储中提取您的数据并将其以批处理形式暴露给训练循环的过程。
Dataset 负责访问和处理单个数据实例。
DataLoader 从 Dataset 中提取数据实例（可以是自动提取或使用您定义的采样器），将它们分批收集，并返回给训练循环使用。 DataLoader 可以与所有类型的数据集一起工作，无论它们包含的数据类型如何。
在本教程中，我们将使用 TorchVision 提供的 Fashion-MNIST 数据集。我们使用 torchvision.transforms.Normalize() 对图像块内容进行零中心化和归一化，并下载训练和验证数据集。
import torch
import torchvision
import torchvision.transforms as transforms
# PyTorch TensorBoard support
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
# Create datasets for training & validation, download if necessary
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform, download=True)
validation_set = torchvision.datasets.FashionMNIST('./data', train=False, transform=transform, download=True)
# Create data loaders for our datasets; shuffle for training, not for validation
training_loader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True)
validation_loader = torch.utils.data.DataLoader(validation_set, batch_size=4, shuffle=False)
# Class labels
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')
# Report split sizes
print('Training set has {} instances'.format(len(training_set)))
print('Validation set has {} instances'.format(len(validation_set)))
像往常一样，让我们可视化数据以进行合理性检查：
import matplotlib.pyplot as plt
import numpy as np
# Helper function for inline image display
def matplotlib_imshow(img, one_channel=False):
if one_channel:
img = img.mean(dim=0)
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
if one_channel:
plt.imshow(npimg, cmap="Greys")
else:
plt.imshow(np.transpose(npimg, (1, 2, 0)))
dataiter = iter(training_loader)
images, labels = next(dataiter)
# Create a grid from the images and show them
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)
print(' '.join(classes[labels[j]] for j in range(4)))
模型 ¶
在本例中，我们将使用 LeNet-5 的一个变体作为模型——如果您观看了本系列之前的视频，应该很熟悉。
import torch.nn as nn
import torch.nn.functional as F
# PyTorch models inherit from torch.nn.Module
class GarmentClassifier(nn.Module):
def __init__(self):
super(GarmentClassifier, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 4 * 4, 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 * 4 * 4)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
model = GarmentClassifier()
损失函数 ¶
在本例中，我们将使用交叉熵损失函数。为了演示目的，我们将创建一些虚拟的输出和标签值批次，将它们通过损失函数，并检查结果。
loss_fn = torch.nn.CrossEntropyLoss()
# NB: Loss functions expect data in batches, so we're creating batches of 4
# Represents the model's confidence in each of the 10 classes for a given input
dummy_outputs = torch.rand(4, 10)
# Represents the correct class among the 10 being tested
dummy_labels = torch.tensor([1, 5, 3, 7])
print(dummy_outputs)
print(dummy_labels)
loss = loss_fn(dummy_outputs, dummy_labels)
print('Total loss for this batch: {}'.format(loss.item()))
优化器 ¶
在这个例子中，我们将使用带有动量的简单随机梯度下降。
尝试对这个优化方案进行一些变体可能会有所帮助：
学习率决定了优化器所采取的步长大小。不同的学习率会对您的训练结果，从准确性和收敛时间方面产生什么影响？
动量会在多个步骤中将优化器推向梯度最强的方向。改变这个值会对您的结果产生什么影响？
尝试一些不同的优化算法，例如平均 SGD、Adagrad 或 Adam。你的结果有何不同？
# Optimizers specified in the torch.optim package
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
训练循环 §
下面，我们有一个执行一个训练周期的函数。它从数据加载器中枚举数据，并在循环的每次迭代中执行以下操作：
从数据加载器中获取一批训练数据
清零优化器的梯度
执行推理 - 即从模型获取输入批次的预测
计算预测集与数据集标签之间的损失
计算学习权重上的反向梯度
告诉优化器执行一次学习步骤——即根据我们选择的优化算法，根据这个批次观察到的梯度调整模型的权重
每隔 1000 个批次报告一次损失。
最后，它报告了最后 1000 个批次的平均每批次损失，以便与验证运行进行比较
def train_one_epoch(epoch_index, tb_writer):
running_loss = 0.
last_loss = 0.
# Here, we use enumerate(training_loader) instead of
# iter(training_loader) so that we can track the batch
# index and do some intra-epoch reporting
for i, data in enumerate(training_loader):
# Every data instance is an input + label pair
inputs, labels = data
# Zero your gradients for every batch!
optimizer.zero_grad()
# Make predictions for this batch
outputs = model(inputs)
# Compute the loss and its gradients
loss = loss_fn(outputs, labels)
loss.backward()
# Adjust learning weights
optimizer.step()
# Gather data and report
running_loss += loss.item()
if i % 1000 == 999:
last_loss = running_loss / 1000 # loss per batch
print(' batch {} loss: {}'.format(i + 1, last_loss))
tb_x = epoch_index * len(training_loader) + i + 1
tb_writer.add_scalar('Loss/train', last_loss, tb_x)
running_loss = 0.
return last_loss
每个 Epoch 的活动
每个 epoch 我们都会想要做几件事情：
通过检查未用于训练的数据集上的相对损失来执行验证，并报告此信息
保存模型副本
在这里，我们将使用 TensorBoard 进行报告。这需要进入命令行启动 TensorBoard，并在另一个浏览器标签页中打开它。
# Initializing in a separate cell so we can easily add more epochs to the same run
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
writer = SummaryWriter('runs/fashion_trainer_{}'.format(timestamp))
epoch_number = 0
EPOCHS = 5
best_vloss = 1_000_000.
for epoch in range(EPOCHS):
print('EPOCH {}:'.format(epoch_number + 1))
# Make sure gradient tracking is on, and do a pass over the data
model.train(True)
avg_loss = train_one_epoch(epoch_number, writer)
running_vloss = 0.0
# Set the model to evaluation mode, disabling dropout and using population
# statistics for batch normalization.
model.eval()
# Disable gradient computation and reduce memory consumption.
with torch.no_grad():
for i, vdata in enumerate(validation_loader):
vinputs, vlabels = vdata
voutputs = model(vinputs)
vloss = loss_fn(voutputs, vlabels)
running_vloss += vloss
avg_vloss = running_vloss / (i + 1)
print('LOSS train {} valid {}'.format(avg_loss, avg_vloss))
# Log the running loss averaged per batch
# for both training and validation
writer.add_scalars('Training vs. Validation Loss',
{ 'Training' : avg_loss, 'Validation' : avg_vloss },
epoch_number + 1)
writer.flush()
# Track best performance, and save the model's state
if avg_vloss < best_vloss:
best_vloss = avg_vloss
model_path = 'model_{}_{}'.format(timestamp, epoch_number)
torch.save(model.state_dict(), model_path)
epoch_number += 1
加载已保存的模型版本：
saved_model = GarmentClassifier()
saved_model.load_state_dict(torch.load(PATH))
加载模型后，它就准备好用于您需要的任何操作了 - 更多训练、推理或分析。
注意，如果您的模型有影响模型结构的构造参数，您需要提供它们，并按保存时的状态配置模型。
其他资源 ¶
PyTorch.org 上的数据工具文档，包括 Dataset 和 DataLoader
关于使用固定内存进行 GPU 训练的说明
TorchVision、TorchText 和 TorchAudio 中可用的数据集文档
PyTorch 中可用的损失函数文档
torch.optim 包的文档，包括优化器和相关工具，如学习率调度
保存和加载模型的详细教程
pytorch.org 的教程部分包含各种训练任务的教程，包括不同领域的分类、生成对抗网络、强化学习等
脚本总运行时间：（0 分钟 0.000 秒）
