快速入门（Quickstart）

本节将介绍机器学习中常见任务的 API。

---

处理数据（Working with Data）

PyTorch 提供了两个处理数据的原语：torch.utils.data.DataLoader 和 torch.utils.data.Dataset。Dataset 存储样本及其对应的标签，DataLoader 在 Dataset 外包装了一个可迭代对象。

import torch                                          # 导入 PyTorch 核心库
from torch import nn                                  # 导入神经网络模块
from torch.utils.data import DataLoader               # 导入数据加载器
from torchvision import datasets                      # 导入视觉数据集模块
from torchvision.transforms import v2                 # 导入数据变换模块 v2 版本

PyTorch 提供了领域特定的库，如 TorchText、TorchVision 和 TorchAudio，这些库都包含数据集。在本教程中，我们将使用 TorchVision 数据集。

torchvision.datasets 模块包含了许多真实世界视觉数据的 Dataset 对象，如 CIFAR、COCO。在本教程中，我们使用 FashionMNIST 数据集。每个 TorchVision Dataset 都包含两个参数：transform 和 target_transform，分别用于修改样本和标签。

# 从开放数据集下载训练数据
training_data = datasets.FashionMNIST(                # 创建 FashionMNIST 训练数据集对象
    root="data",                                      # 数据存储路径
    train=True,                                       # 指定为训练集
    download=True,                                    # 如果数据不存在则从网络下载
    transform=v2.Compose([                            # 对图像进行变换组合
        v2.ToImage(),                                 # 将数据转为 torchvision 的 Image 张量格式
        v2.ToDtype(torch.float32, scale=True)         # 转为 float32 类型，并将像素值缩放到 [0,1]
    ]),
)

# 从开放数据集下载测试数据
test_data = datasets.FashionMNIST(                    # 创建 FashionMNIST 测试数据集对象
    root="data",                                      # 数据存储路径
    train=False,                                      # 指定为测试集
    download=True,                                    # 如果数据不存在则从网络下载
    transform=v2.Compose([                            # 对图像进行变换组合（与训练集相同）
        v2.ToImage(),                                 # 将数据转为 torchvision 的 Image 张量格式
        v2.ToDtype(torch.float32, scale=True)         # 转为 float32 类型，并将像素值缩放到 [0,1]
    ]),
)

我们将 Dataset 作为参数传递给 DataLoader。这将在数据集外包装一个可迭代对象，并支持自动批处理、采样、打乱和多进程数据加载。这里我们定义批大小为 64，即 dataloader 可迭代对象中的每个元素将返回一个包含 64 个特征和标签的批次。

batch_size = 64                                       # 定义批大小为 64

# 创建数据加载器
train_dataloader = DataLoader(training_data, batch_size=batch_size)  # 用训练数据集创建数据加载器
test_dataloader = DataLoader(test_data, batch_size=batch_size)       # 用测试数据集创建数据加载器

for X, y in test_dataloader:                          # 遍历测试数据加载器
    print(f"Shape of X [N, C, H, W]: {X.shape}")      # 打印输入特征的形状：批次大小、通道数、高度、宽度
    print(f"Shape of y: {y.shape} {y.dtype}")         # 打印标签的形状和数据类型
    break                                             # 只查看第一个批次后跳出循环

输出：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64

---

创建模型（Creating Models）

要在 PyTorch 中定义神经网络，我们创建一个继承自 nn.Module 的类。我们在 __init__ 函数中定义网络的层，并在 forward 函数中指定数据如何通过网络。为了加速神经网络中的运算，我们将其移动到加速器上，如 CUDA、MPS、MTIA 或 XPU。如果当前加速器可用，我们将使用它；否则，使用 CPU。

# 检查加速器是否可用，如果可用则使用加速器，否则使用 CPU
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")                       # 打印当前使用的设备类型

# 定义模型
class NeuralNetwork(nn.Module):                       # 定义神经网络类，继承自 nn.Module
    def __init__(self):                               # 初始化方法
        super().__init__()                            # 调用父类的初始化方法
        self.flatten = nn.Flatten()                   # 创建展平层，将 2D 图像展平为 1D 向量
        self.linear_relu_stack = nn.Sequential(       # 创建顺序容器，按顺序包含各层
            nn.Linear(28*28, 512),                    # 第一个全连接层：输入 784 维，输出 512 维
            nn.ReLU(),                                # ReLU 激活函数，引入非线性
            nn.Linear(512, 512),                      # 第二个全连接层：输入 512 维，输出 512 维
            nn.ReLU(),                                # ReLU 激活函数
            nn.Linear(512, 10)                        # 第三个全连接层：输入 512 维，输出 10 维（对应 10 个类别）
        )

    def forward(self, x):                             # 前向传播方法，定义数据如何通过网络
        x = self.flatten(x)                           # 将输入图像展平
        logits = self.linear_relu_stack(x)            # 通过全连接层序列得到原始预测值（logits）
        return logits                                 # 返回预测值

model = NeuralNetwork().to(device)                    # 创建模型实例并移动到指定设备
print(model)                                          # 打印模型结构

输出：

Using cuda device
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

---

优化模型参数（Optimizing the Model Parameters）

要训练模型，我们需要一个损失函数和一个优化器。

loss_fn = nn.CrossEntropyLoss()                       # 定义交叉熵损失函数，适用于多分类任务
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)  # 定义 SGD 随机梯度下降优化器，学习率为 0.001

在单个训练循环中，模型对训练数据集进行预测（以批次形式输入），并通过反向传播预测误差来调整模型的参数。

def train(dataloader, model, loss_fn, optimizer):     # 定义训练函数
    size = len(dataloader.dataset)                    # 获取数据集的样本总数
    model.train()                                     # 将模型设置为训练模式
    for batch, (X, y) in enumerate(dataloader):       # 遍历数据加载器中的每个批次
        X, y = X.to(device), y.to(device)             # 将数据和标签移动到指定设备

        # 计算预测误差
        pred = model(X)                               # 前向传播，得到模型预测
        loss = loss_fn(pred, y)                       # 计算预测值与真实标签之间的损失

        # 反向传播
        loss.backward()                               # 反向传播，计算梯度
        optimizer.step()                              # 根据梯度更新模型参数
        optimizer.zero_grad()                         # 清零梯度，防止梯度累积

        if batch % 100 == 0:                          # 每 100 个批次打印一次
            loss, current = loss.item(), (batch + 1) * len(X)  # 获取当前损失值和已处理的样本数
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")  # 打印损失和进度

我们还在测试数据集上检查模型的性能，以确保它正在学习。

def test(dataloader, model, loss_fn):                 # 定义测试函数
    size = len(dataloader.dataset)                    # 获取数据集的样本总数
    num_batches = len(dataloader)                     # 获取批次总数
    model.eval()                                      # 将模型设置为评估模式
    test_loss, correct = 0, 0                         # 初始化测试损失和正确预测计数
    with torch.no_grad():                             # 禁用梯度计算，节省内存和计算
        for X, y in dataloader:                       # 遍历测试数据加载器
            X, y = X.to(device), y.to(device)         # 将数据和标签移动到指定设备
            pred = model(X)                           # 前向传播，得到模型预测
            test_loss += loss_fn(pred, y).item()      # 累加批次损失
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  # 统计正确预测数
    test_loss /= num_batches                          # 计算平均测试损失
    correct /= size                                   # 计算准确率
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")  # 打印测试结果

训练过程在若干次迭代（epoch）中进行。在每个 epoch 中，模型学习参数以做出更好的预测。我们在每个 epoch 打印模型的准确率和损失；我们希望看到准确率随着每个 epoch 增加而损失减少。

epochs = 5                                            # 训练 5 个 epoch
for t in range(epochs):                               # 遍历每个 epoch
    print(f"Epoch {t+1}\n-------------------------------")  # 打印当前 epoch 编号
    train(train_dataloader, model, loss_fn, optimizer)  # 执行训练循环
    test(test_dataloader, model, loss_fn)              # 执行测试循环
print("Done!")                                        # 训练完成

输出：

Epoch 1
-------------------------------
loss: 2.302590  [   64/60000]
loss: 2.291843  [ 6464/60000]
loss: 2.274084  [12864/60000]
loss: 2.267996  [19264/60000]
loss: 2.248032  [25664/60000]
loss: 2.235034  [32064/60000]
loss: 2.235448  [38464/60000]
loss: 2.206563  [44864/60000]
loss: 2.203167  [51264/60000]
loss: 2.173855  [57664/60000]
Test Error:
 Accuracy: 48.5%, Avg loss: 2.165248

Epoch 2
-------------------------------
loss: 2.171659  [   64/60000]
loss: 2.168588  [ 6464/60000]
loss: 2.110749  [12864/60000]
loss: 2.124941  [19264/60000]
loss: 2.070731  [25664/60000]
loss: 2.026289  [32064/60000]
loss: 2.046067  [38464/60000]
loss: 1.972416  [44864/60000]
loss: 1.973158  [51264/60000]
loss: 1.901610  [57664/60000]
Test Error:
 Accuracy: 59.3%, Avg loss: 1.900692

Epoch 3
-------------------------------
loss: 1.927220  [   64/60000]
loss: 1.909643  [ 6464/60000]
loss: 1.786996  [12864/60000]
loss: 1.823585  [19264/60000]
loss: 1.710886  [25664/60000]
loss: 1.668456  [32064/60000]
loss: 1.684839  [38464/60000]
loss: 1.584657  [44864/60000]
loss: 1.603412  [51264/60000]
loss: 1.505005  [57664/60000]
Test Error:
 Accuracy: 60.7%, Avg loss: 1.522182

Epoch 4
-------------------------------
loss: 1.580342  [   64/60000]
loss: 1.556963  [ 6464/60000]
loss: 1.401661  [12864/60000]
loss: 1.472037  [19264/60000]
loss: 1.356207  [25664/60000]
loss: 1.352045  [32064/60000]
loss: 1.366575  [38464/60000]
loss: 1.284744  [44864/60000]
loss: 1.322070  [51264/60000]
loss: 1.234389  [57664/60000]
Test Error:
 Accuracy: 62.9%, Avg loss: 1.253126

Epoch 5
-------------------------------
loss: 1.323655  [   64/60000]
loss: 1.310964  [ 6464/60000]
loss: 1.143728  [12864/60000]
loss: 1.248905  [19264/60000]
loss: 1.131093  [25664/60000]
loss: 1.151796  [32064/60000]
loss: 1.175902  [38464/60000]
loss: 1.103851  [44864/60000]
loss: 1.148382  [51264/60000]
loss: 1.078236  [57664/60000]
Test Error:
 Accuracy: 64.5%, Avg loss: 1.088746

Done!

---

保存模型（Saving Models）

保存模型的常用方法是序列化内部状态字典（包含模型参数）。

torch.save(model.state_dict(), "model.pth")          # 将模型参数（状态字典）保存到文件 model.pth
print("Saved PyTorch Model State to model.pth")       # 打印保存成功提示

输出：

Saved PyTorch Model State to model.pth

加载模型（Loading Models）

加载模型的过程包括重新创建模型结构并将状态字典加载到其中。

model = NeuralNetwork().to(device)                    # 重新创建模型实例并移动到设备
model.load_state_dict(torch.load("model.pth", weights_only=True))  # 从文件加载模型参数，weights_only=True 是安全最佳实践

输出：

<All keys matched successfully>

此模型现在可用于进行预测。

classes = [                                           # 定义类别名称列表
    "T-shirt/top",                                    # 0: T恤
    "Trouser",                                        # 1: 裤子
    "Pullover",                                       # 2: 套衫
    "Dress",                                          # 3: 连衣裙
    "Coat",                                           # 4: 外套
    "Sandal",                                         # 5: 凉鞋
    "Shirt",                                          # 6: 衬衫
    "Sneaker",                                        # 7: 运动鞋
    "Bag",                                            # 8: 包
    "Ankle boot",                                     # 9: 踝靴
]

model.eval()                                          # 将模型设置为评估模式
x, y = test_data[0][0], test_data[0][1]               # 获取测试数据集的第一张图像和标签
with torch.no_grad():                                 # 禁用梯度计算
    x = x.to(device)                                  # 将图像移动到设备
    pred = model(x)                                   # 模型预测
    predicted, actual = classes[pred[0].argmax(0)], classes[y]  # 获取预测类别和真实类别
    print(f'Predicted: "{predicted}", Actual: "{actual}"')     # 打印预测结果和真实标签

输出：

Predicted: "Ankle boot", Actual: "Ankle boot"