PyTorch 入门指南

PyTorch 是当今学术界与工业界最主流的深度学习框架——动态图、自动微分、Pythonic 的 API 让”定义网络、训练模型”变得直觉而灵活。本页带你从 Tensor 到完整训练循环一次走通。前置阅读：NumPy 科学计算入门、反向传播。

这个库是什么

PyTorch 是 Meta（原 Facebook）AI 研究院开源的深度学习框架，2017 年发布。它围绕三个核心能力构建：多维张量（Tensor，类似 NumPy 但支持 GPU）、自动微分（autograd，即自动计算梯度的系统）、以及构建模型的模块化系统（nn.Module）。

它之所以击败 TensorFlow 成为研究界首选，关键在于动态计算图（define-by-run，即”运行时即时构建计算图”的模式）：网络的计算图在每次前向传播时即时构建，意味着你可以用熟悉的 if、for、print 调试神经网络，就像写普通 Python 代码。打个比方，TensorFlow 1.x 像是”先画完整张蓝图再施工”，而 PyTorch 是”边设计边施工”，随时能改。如今 GPT、LLaMA、Stable Diffusion 等主流大模型的官方实现都基于 PyTorch。

截至 2025 年，PyTorch 已迭代至 2.x 系列（当前稳定版 2.13），核心趋势从”灵活优先”逐步走向”灵活与性能兼得”——2.0 引入的 torch.compile（即时编译加速器，能把 Python 动态代码编译成优化的 GPU kernel）让 PyTorch 在不牺牲易用性的前提下大幅缩小了与 TensorFlow 静态图的性能差距，而 2025 年新增的 FlexAttention（灵活注意力机制，允许用少量代码定义任意注意力模式并自动获得 FlashAttention 级别的加速）、FSDP2（第二代完全分片数据并行，用于在多 GPU 上训练超大模型）等功能进一步巩固了它在大模型训练领域的统治地位。

安装与环境配置

# 方式一：CPU 版（pip）
pip install torch torchvision

# 方式二：GPU 版（CUDA 12.6，按官网指引选择对应版本）
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126

# 验证安装与 GPU 可用性
python -c "import torch; print(torch.__version__); print('CUDA:', torch.cuda.is_available())"

PyTorch 安装包较大（GPU 版数 GB），建议按 pytorch.org 首页的交互式选择器生成准确命令。导入约定写为 import torch，视觉模块为 import torchvision。

版本提示（2025）：从 PyTorch 2.11 起，PyPI 上的默认 wheel 已升级到 CUDA 13.0。如果你的 GPU 驱动较旧（仅支持 CUDA 12.x），需要显式指定 --index-url .../cu126 或 cu128。Maxwell / Pascal / Volta 等老架构 GPU 请使用 CUDA 12.6 构建。Python 3.14 已获全面支持，Python 3.15 wheel 也已在 2.13 中提供。

注意 PyTorch 2.x 已内置编译后端 torch.compile，能显著加速训练而无需改代码：

import torch

model = MyModel()

# 一行代码加速：torch.compile 会把模型编译成优化后的计算图
# 首次调用有编译开销，后续运行会快很多（尤其是 Transformer 等计算密集型模型）
compiled_model = torch.compile(model)

核心概念

Tensor：会算梯度、能上 GPU 的 NumPy 数组

Tensor 是 PyTorch 的基本数据结构，你可以把它理解为”超级 NumPy 数组”——它具备 ndarray 的所有运算能力（广播、切片、矩阵乘法），但额外支持两件事：自动记录运算以计算梯度，以及无缝迁移到 GPU。

创建：torch.tensor([1,2,3])、torch.zeros(3,4)、torch.randn(100,10)。
设备迁移：.to("cuda") 把数据搬到 GPU，运算随即在 GPU 上进行。
与 NumPy 互转：tensor.numpy() 与 torch.from_numpy(arr) 共享内存，零拷贝。

为什么需要 Tensor 而不直接用 Python 列表或 NumPy？ 神经网络的核心运算是大规模矩阵乘法。一个含 100 万参数的模型，每次前向传播要做上亿次浮点运算。Tensor 底层用 C++/CUDA 实现，能在 GPU 上同时启动数千个计算核心并行处理，比 Python 循环快好几个数量级。用个比喻：NumPy 是一把高效的手术刀，Tensor 则是一间设备齐全的手术室——不仅有手术刀，还有麻醉师（autograd）和护士团队（GPU 并行）。

autograd：自动微分引擎

训练神经网络的核心是反向传播（用链式法则算梯度），手写梯度极易出错。autograd（自动微分引擎，即自动计算梯度的系统）在前向传播时自动记录所有运算，构建一张计算图（computational graph，一种有向无环图 DAG，记录运算之间的依赖关系）；调用 .backward() 时，它从输出反向遍历这张图，自动算出每个参数的梯度。你只需写前向逻辑，梯度交给引擎——这就像自动驾驶：你只管设定目的地（loss），引擎自动算出该怎么打方向盘（梯度）。

数学原理：链式法则的自动化

autograd 的核心是微积分中的链式法则（chain rule）。假设有一个复合函数 $y = f(g(x))$ ，那么：

$\frac{dy}{dx} = \frac{dy}{dg} \cdot \frac{dg}{dx}$

在神经网络中，前向传播就是一个嵌套的复合函数。以三层网络为例，输出 $\hat{y}$ 可以写成：

$\hat{y} = f_3(W_3 \cdot f_2(W_2 \cdot f_1(W_1 \cdot x + b_1) + b_2) + b_3)$

要更新参数 $W_1$ ，需要算 $\frac{\partial L}{\partial W_1}$ （损失函数对 $W_1$ 的梯度）。按链式法则：

$\frac{\partial L}{\partial W_1} = \frac{\partial L}{\partial \hat{y}} \cdot \frac{\partial \hat{y}}{\partial h_2} \cdot \frac{\partial h_2}{\partial h_1} \cdot \frac{\partial h_1}{\partial W_1}$

其中 $h_1, h_2$ 是各隐藏层的输出。手写这个公式又长又容易出错，而 autograd 把每一步运算的局部导数（Jacobian）记录下来，反向传播时逐个相乘，自动得到最终梯度。

具体例子： $y = x^2 + 3x + 1$ ，求 $x=2$ 时的 $\frac{dy}{dx}$ 。

前向： $y = 4 + 6 + 1 = 11$
反向： $\frac{dy}{dx} = 2x + 3 = 2(2) + 3 = 7$

autograd 就是自动帮你做这件事——你只管写 y = x**2 + 3*x + 1，它自动算出梯度是 7。

nn.Module：搭积木式构建网络

nn.Module 是所有网络层和模型的基类。定义网络只需继承它、在 __init__ 中声明各层、在 forward 中描述数据如何流过这些层。nn.Linear（全连接层，即 $y = Wx + b$ ）、nn.Conv2d（卷积层）、nn.Transformer 等内置层都是积木，而 nn.Sequential 则像一个穿串器，把多个层按顺序串起来。

nn.Module 的精妙之处在于：每个 Module 会自动追踪自己的可训练参数（parameters()），你调用 model.parameters() 就能拿到所有需要优化的权重，直接传给优化器。这种”模块即参数容器”的设计让复杂网络的参数管理变得透明。

架构与工作流

动手实践

例 1：Tensor 创建与运算

import torch

# 创建张量（类似 NumPy）
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.randn(3, 4)              # 3x4 标准正态随机矩阵
zeros = torch.zeros(2, 3)

# 运算与广播——和 NumPy 几乎完全一致
print(a + 10)                      # [11, 12, 13] 标量广播
print(b @ b.T)                     # 矩阵乘法，结果 3x3
print(b.shape, b.dtype)            # torch.Size([3, 4]) torch.float32

# 与 NumPy 互转共享内存
import numpy as np
arr = np.ones(5)
t = torch.from_numpy(arr)          # 零拷贝转换
arr[0] = 99
print(t[0].item())                 # 99.0，修改互相影响

例 2：自动微分

import torch

# requires_grad=True 让 autograd 追踪这个张量的运算
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1             # y = x² + 3x + 1，当 x=2 时 y=11
y.backward()                       # 自动反向传播，计算 dy/dx
print(x.grad)                      # 2x + 3 = 7.0 —— 这就是解析导数

# 训练神经网络时，梯度是 autograd 自动算的，你只需写前向 loss
w = torch.randn(10, 1, requires_grad=True)   # 可训练参数
loss = (w ** 2).sum()                          # 假设的损失
loss.backward()
print(w.grad.shape)                            # torch.Size([10, 1])

多变量梯度怎么算？ autograd 不仅能算标量对向量的梯度，还能处理更复杂的情况。看下面这个多变量例子：

import torch

# 模拟一个简单的线性回归：y_hat = w·x + b
# 损失函数 L = (y_hat - y)² ，求 L 对 w 和 b 的偏导
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=False)   # 输入
y = torch.tensor([2.0, 4.0, 6.0], requires_grad=False)   # 真实标签（恰好 y = 2x）

w = torch.tensor(1.0, requires_grad=True)   # 初始权重
b = torch.tensor(0.0, requires_grad=True)   # 初始偏置

y_hat = w * x + b                            # 前向：预测值 [1, 2, 3]
loss = ((y_hat - y) ** 2).mean()             # MSE 损失 = mean([1, 4, 9]) = 4.667

loss.backward()
print(f"w.grad = {w.grad:.4f}")   # -4.667（解析：dL/dw = (2/3)·Σ(w·xi+b-yi)·xi）
print(f"b.grad = {b.grad:.4f}")   # -2.000（解析：dL/db = (2/3)·Σ(w·xi+b-yi)）
# 梯度为负，说明 w 和 b 应该增大，朝 y=2x 的方向调整

例 3：用 nn.Module 定义模型

import torch
import torch.nn as nn

# 定义一个简单的两层全连接网络
class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),   # 第一层
            nn.ReLU(),                   # 激活函数
            nn.Linear(hidden, out_dim),  # 输出层
        )

    def forward(self, x):               # 描述前向计算
        return self.net(x)

model = MLP(20, 64, 3)                  # 20 维输入，3 类输出
print(model)                            # 打印网络结构
x = torch.randn(8, 20)                  # 8 个样本，每个 20 维
print(model(x).shape)                    # torch.Size([8, 3])

nn.Linear 在数学上做了什么？ 它就是矩阵乘法加偏置： $\mathbf{y} = \mathbf{W}\mathbf{x} + \mathbf{b}$ 。其中 $\mathbf{W}$ 的形状是 (out_features, in_features)， $\mathbf{b}$ 的形状是 (out_features,)。nn.Linear(20, 64) 会自动创建一个 64×20 的权重矩阵和一个 64 维偏置向量，共 $64 \times 20 + 64 = 1344$ 个参数。你可以用 sum(p.numel() for p in model.parameters()) 数出模型总参数量。

例 4：完整训练循环

import torch
import torch.nn as nn

# 准备随机数据：1000 个样本，20 维特征，3 类
X = torch.randn(1000, 20)
y = torch.randint(0, 3, (1000,))

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3))
criterion = nn.CrossEntropyLoss()                  # 多分类交叉熵
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

# 标准训练循环：前向 → 算 loss → 反向 → 更新
for epoch in range(50):
    pred = model(X)                                # 前向传播
    loss = criterion(pred, y)                      # 计算损失

    optimizer.zero_grad()                          # 清空旧梯度（关键！）
    loss.backward()                                # 反向传播算梯度
    optimizer.step()                               # 用梯度更新参数

    if epoch % 10 == 0:
        acc = (pred.argmax(1) == y).float().mean()
        print(f"epoch {epoch}, loss={loss.item():.4f}, acc={acc:.3f}")

交叉熵损失的数学推导

上面用到的 CrossEntropyLoss（交叉熵损失，衡量预测概率分布与真实分布差异的指标）是分类任务最常用的损失函数。对于 $C$ 类分类问题，给定样本的真实类别 $c$ 和模型的原始输出 $\mathbf{z} = (z_1, z_2, \ldots, z_C)$ （称为 logits，即未经 softmax 归一化的原始得分），交叉熵损失为：

$L = -\log \frac{e^{z_c}}{\sum_{j=1}^{C} e^{z_j}} = -z_c + \log \sum_{j=1}^{C} e^{z_j}$

第一步用了 softmax 把 logits 转成概率： $p_c = \frac{e^{z_c}}{\sum_j e^{z_j}}$ ，然后取 $-\log p_c$ 。

它对 logits 的梯度非常简洁：

$\frac{\partial L}{\partial z_j} = p_j - \mathbb{1}[j = c]$

其中 $p_j$ 是 softmax 后第 $j$ 类的概率， $\mathbb{1}[j=c]$ 在 $j$ 等于真实类别时为 1，否则为 0。直觉理解：梯度就是”预测概率减去真实标签（one-hot）“——如果模型对正确类别预测概率为 1，梯度为 0（完美预测）；概率越低，梯度越大（推力越强）。PyTorch 的 CrossEntropyLoss 内部已包含 softmax，所以网络最后一层直接输出 logits 即可。

实战：从零搭建 CNN 图像分类器

上面四个例子用随机数据演示了 PyTorch 的基本套路。本节把所有部件拼起来，搭建一个完整可运行的卷积神经网络（CNN），在 MNIST 手写数字数据集上做 10 分类。这是本页最重要的部分——看完它，你就掌握了”定义网络 → 喂数据 → 训练 → 评估 → 保存”的完整闭环，足以应对绝大多数入门级图像分类任务。

第 1 步：数据准备

深度学习的第一步永远是准备数据。torchvision 内置了常见数据集的下载与加载工具，配合 DataLoader（数据加载器，负责把数据集切成批次、打乱顺序、多线程读取的工具）实现自动分批、打乱、多进程加载。

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 定义数据预处理：把 PIL 图片转成 Tensor 并归一化到 [0,1] 区间
# ToTensor() 会做两件事：(1) 把 H×W×C 的 PIL 图片转成 C×H×W 的 Tensor；(2) 把像素值从 [0,255] 缩放到 [0,1]
transform = transforms.Compose([
    transforms.ToTensor(),
])

# 下载并加载 MNIST（手写数字 0-9，28×28 灰度图，10 个类别）
# FashionMNIST 换成 datasets.FashionMNIST 即可，接口完全一样
# CIFAR-10 换成 datasets.CIFAR10(root=..., download=True)，但图片是 3 通道彩色
train_dataset = datasets.MNIST(
    root="./data",        # 数据存放目录
    train=True,           # 加载训练集
    download=True,        # 首次运行自动下载
    transform=transform,  # 应用上面定义的预处理
)
test_dataset = datasets.MNIST(
    root="./data",
    train=False,          # 加载测试集
    download=True,
    transform=transform,
)

# DataLoader：把数据集切成一个一个 batch，自动打乱、多进程加速读取
train_loader = DataLoader(
    train_dataset,
    batch_size=64,        # 每批 64 张图片
    shuffle=True,         # 每个 epoch 重新打乱顺序，防止模型记住数据顺序
    num_workers=2,        # 用 2 个子进程并行加载数据，不阻塞训练主循环
)
test_loader = DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False,        # 测试集不需要打乱，方便对比结果
    num_workers=2,
)

# 取出一个 batch 看看数据形状
images, labels = next(iter(train_loader))
print(images.shape)   # torch.Size([64, 1, 28, 28])
print(labels.shape)   # torch.Size([64])
print(labels[:8])     # tensor([3, 0, 4, 1, 7, 2, 9, 5]) —— 每张图对应的数字标签

理解数据维度——这是后续一切的基础。 PyTorch 中图像数据的形状始终是 (B, C, H, W) 四个维度：

B（Batch）：一次同时处理多少张图。上面设的 batch_size=64，所以 B=64。批量处理能充分利用 GPU 并行算力，也让梯度更稳定。
C（Channel）：图像的通道数。MNIST 是灰度图，C=1；CIFAR-10 / 自然照片是 RGB 彩色图，C=3。注意 PyTorch 把通道放在第二维（C×H×W），而 PIL/OpenCV 默认是 H×W×C，ToTensor() 会自动帮你调换。
H、W（Height、Width）：图像的高和宽（像素数）。MNIST 图片是 28×28，所以 H=28、W=28。

所以 images.shape = (64, 1, 28, 28) 的含义就是：这批数据有 64 张图，每张 1 个通道（灰度），高 28 像素、宽 28 像素。后续设计网络时，每一层的输入输出维度都要和这个形状对上。

第 2 步：设计 CNN 网络（核心——维度对应）

下面定义一个两层卷积的 CNN 分类器。每一层都标注了输入输出维度，请务必跟着维度变化读一遍——理解维度流转，就理解了 CNN 的设计逻辑。

import torch.nn as nn

class MNISTClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # 第一个卷积块：卷积 + 激活 + 池化
        self.conv1 = nn.Conv2d(
            in_channels=1,        # 输入通道数：MNIST 灰度图是 1
            out_channels=32,      # 输出通道数：提取 32 种特征
            kernel_size=3,        # 3×3 卷积核
            padding=1,            # 周围补 1 圈零，保持尺寸不变
        )
        self.relu1 = nn.ReLU()    # 激活函数，引入非线性
        self.pool1 = nn.MaxPool2d(2)   # 2×2 最大池化，尺寸减半

        # 第二个卷积块
        self.conv2 = nn.Conv2d(
            in_channels=32,       # 上一步输出了 32 个通道，这里输入就是 32
            out_channels=64,      # 提取 64 种更高层特征
            kernel_size=3,
            padding=1,
        )
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)

        # 全连接分类头
        self.flatten = nn.Flatten()           # 把 (B,64,7,7) 展平成 (B, 3136)
        self.fc1 = nn.Linear(64 * 7 * 7, 128) # 3136 维 → 128 维
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)  # 128 维 → 10 类

    def forward(self, x):
        # x 形状变化（以 batch_size=64 为例）：
        x = self.conv1(x)   # (64, 1, 28, 28) → (64, 32, 28, 28)
        x = self.relu1(x)   # 激活不改形状：(64, 32, 28, 28)
        x = self.pool1(x)   # 尺寸减半：(64, 32, 28, 28) → (64, 32, 14, 14)

        x = self.conv2(x)   # (64, 32, 14, 14) → (64, 64, 14, 14)
        x = self.relu2(x)   # (64, 64, 14, 14)
        x = self.pool2(x)   # (64, 64, 14, 14) → (64, 64, 7, 7)

        x = self.flatten(x)  # (64, 64, 7, 7) → (64, 3136)，把三维特征图拉成一维向量
        x = self.fc1(x)      # (64, 3136) → (64, 128)
        x = self.relu3(x)    # (64, 128)
        x = self.fc2(x)      # (64, 128) → (64, 10)，输出 10 个类别的得分
        return x

把整张数据流图画出来，维度变化一目了然：

几个关键设计原则：

padding=1 配合 kernel_size=3：卷积后尺寸不变。公式是 $H_{\text{out}} = \lfloor \frac{H_{\text{in}} + 2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} \rfloor + 1$ ，代入 $\frac{28 + 2 - 3}{1} + 1 = 28$ ，所以 28×28 卷积后还是 28×28，尺寸缩小只发生在池化层。
MaxPool2d(2)：2×2 窗口取最大值，高和宽各减半。它没有可学习参数，只做下采样，既减少计算量又扩大感受野（receptive field，即一个输出像素能”看到”的输入区域大小）。
in_channels 必须等于上一层的 out_channels：第一层卷积 1→32，第二层就必须写 32→64。通道数不匹配会直接报错。
Flatten 后的维度 = 通道数 × 高 × 宽：两次池化后特征图是 64×7×7 = 3136，所以 fc1 的输入维度是 3136。算错这个数是最常见的维度报错来源。
输出层不接激活函数：CrossEntropyLoss 内部已包含 softmax，所以网络的最后一层直接输出原始得分（logits）即可。

卷积运算的数学直觉

卷积核（kernel / filter）是一个小的权重矩阵（如 3×3），它在输入图像上滑动，每到一个位置就和覆盖的区域做逐元素乘法再求和，生成输出的一个像素。用公式表达，对于输入特征图 $X$ 和卷积核 $K$ （大小 $k \times k$ ）：

$Y[i, j] = \sum_{m=0}^{k-1} \sum_{n=0}^{k-1} X[i+m, j+n] \cdot K[m, n] + b$

其中 $b$ 是偏置。直观理解：卷积核就像一个”特征探测器”——如果图像某区域的纹理和核的模式匹配，乘积之和就大，输出值就高。一个 Conv2d(1, 32, 3) 层有 32 个不同的探测器，分别寻找边缘、角点、纹理等低级特征。

用比喻来理解：想象你看一张风景照，卷积核就是你的”眼睛扫描策略”——你不会一次性看完全部细节，而是用一个小窗口逐块扫过画面，每块提取一个特征（“这里有边缘""这里有颜色渐变”）。卷积层做了同样的事，只不过”窗口”里的权重是学习出来的。

第 3 步：训练循环（完整代码，逐行注释）

现在把网络、损失函数、优化器组装起来，跑训练循环。这是深度学习最核心的”五步咒语”：zero_grad → 前向 → 算 loss → 反向 → 更新。

import torch

# 自动选择设备：有 GPU 就用 GPU，没有就用 CPU
# 在 Mac 上还可以用 torch.device("mps") 调用 Apple 芯片的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MNISTClassifier(num_classes=10).to(device)   # 模型搬到设备上
criterion = nn.CrossEntropyLoss()                     # 多分类交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)  # Adam 优化器

num_epochs = 5   # 遍历整个数据集的次数

for epoch in range(num_epochs):
    model.train()   # 切换到训练模式（影响 Dropout / BatchNorm 的行为）
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in train_loader:
        # 把数据搬到和模型相同的设备上（GPU 上训练就必须这一步）
        images = images.to(device)   # (B, 1, 28, 28)
        labels = labels.to(device)   # (B,)

        # —— 标准训练五步 ——
        optimizer.zero_grad()          # 1. 清空上一步的残余梯度（PyTorch 梯度默认累积，不清零会出错）
        outputs = model(images)        # 2. 前向传播：把图片喂给网络，得到预测得分 (B, 10)
        loss = criterion(outputs, labels)  # 3. 计算损失：预测得分与真实标签之间的差距
        loss.backward()                # 4. 反向传播：autograd 自动计算每个参数的梯度
        optimizer.step()               # 5. 更新参数：优化器用梯度调整权重，朝 loss 减小的方向走一步

        # 统计本 batch 的指标
        running_loss += loss.item() * images.size(0)   # 累加损失（乘以 batch 大小用于求平均）
        predicted = outputs.argmax(dim=1)              # 取得分最大的类别作为预测结果
        correct += (predicted == labels).sum().item()  # 统计预测正确的数量
        total += labels.size(0)

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    print(f"Epoch [{epoch+1}/{num_epochs}]  loss={epoch_loss:.4f}  acc={epoch_acc:.4f}")

运行后你会看到类似这样的输出，loss 逐轮下降、准确率逐轮上升：

Epoch [1/5]  loss=0.1423  acc=0.9561
Epoch [2/5]  loss=0.0468  acc=0.9854
Epoch [3/5]  loss=0.0321  acc=0.9898
Epoch [4/5]  loss=0.0235  acc=0.9924
Epoch [5/5]  loss=0.0187  acc=0.9940

训练五步各自在干什么？

optimizer.zero_grad()：PyTorch 的梯度默认是累加的（为了支持梯度累加等技巧）。如果不清零，这一轮的梯度会叠到上一轮上，训练直接崩溃。这是新手最容易忘的一行。
outputs = model(images)：前向传播，数据从输入层逐层流到输出层，得到预测值。这一步同时构建了计算图（记录了哪些参数参与了哪些运算）。
loss = criterion(...)：用一个标量衡量”预测有多偏离真相”。分类任务用交叉熵，回归任务用均方误差。这个标量就是反向传播的起点。
loss.backward()：从 loss 出发，沿计算图反向走，用链式法则算出每个参数对 loss 的影响（梯度）。梯度存在各参数的 .grad 属性里。
optimizer.step()：优化器根据算好的梯度，按特定策略（SGD、Adam 等）更新参数。这才是模型真正”学习”的瞬间——权重变了，下一轮预测就会更准。

Adam 优化器的数学直觉

上面用的 Adam（Adaptive Moment Estimation，自适应矩估计优化器）是目前最常用的优化算法。它结合了动量法和自适应学习率两个思路：

一阶矩估计（动量）： $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ —— 梯度的指数移动平均，平滑噪声。
二阶矩估计（自适应学习率）： $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$ —— 梯度平方的移动平均，估计每个参数的梯度波动幅度。
参数更新： $\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t$ —— 学习率除以 $\sqrt{\hat{v}_t}$ ，意味着梯度波动大的参数自动缩小步长，波动小的参数保持步长。

直觉理解：Adam 像一个有记忆的登山者——它不仅记住之前走过的方向（动量），还根据每个方向的”路况”自动调整步长（自适应）。这使得它对学习率不太敏感，通常 lr=1e-3 就能工作得很好，是新手的首选优化器。

第 4 步：评估与推理

训练完的模型要放到测试集上评估真实性能，并学会对单张图片做推理。

# —— 在测试集上评估整体准确率 ——
model.eval()   # 切换到评估模式：关闭 Dropout、冻结 BatchNorm
test_correct = 0
test_total = 0

with torch.no_grad():   # 推理时不需要计算梯度，省显存、加速
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)                  # 前向传播，只做预测不反向
        predicted = outputs.argmax(dim=1)        # 取得分最高的类别
        test_correct += (predicted == labels).sum().item()
        test_total += labels.size(0)

test_acc = test_correct / test_total
print(f"测试集准确率: {test_acc:.4f}")
# 输出示例：测试集准确率: 0.9892


# —— 对单张图片做推理 ——
import matplotlib.pyplot as plt

# 从测试集取一张图（也可以用 PIL/OpenCV 读自己的图片，再用 transform 处理）
image, true_label = test_dataset[0]      # image 形状 (1, 28, 28)，值域 [0,1]

# 模型期望的输入是 (B, C, H, W)，单张图要补一个 batch 维 → (1, 1, 28, 28)
input_batch = image.unsqueeze(0).to(device)

model.eval()
with torch.no_grad():
    output = model(input_batch)          # (1, 10)，10 个类别的得分
    predicted_label = output.argmax(dim=1).item()

print(f"真实标签: {true_label}, 预测结果: {predicted_label}")
# 可视化（可选）
plt.imshow(image.squeeze(), cmap="gray")
plt.title(f"Predicted: {predicted_label}")
plt.axis("off")
plt.show()

MNIST 推理可视化

两个关键点：model.eval() 改变模型行为（让 Dropout/BatchNorm 进入推理状态），torch.no_grad() 关闭梯度计算（省内存）。评估和推理时两者都要用。

第 5 步：GPU 加速

上面的代码里已经用到了 GPU，这里集中说明几个要点。

# 判断是否有 GPU
print(torch.cuda.is_available())   # True 表示有 NVIDIA GPU 且装了 CUDA 版 PyTorch

# 创建设备对象
device = torch.device("cuda")   # 或 "cpu"，Mac 上可用 "mps"

# 两样东西都必须搬到 GPU，缺一不可：
model = model.to(device)                    # 模型的参数搬到 GPU
images = images.to(device)                  # 数据搬到 GPU
labels = labels.to(device)

# 之后所有运算（model(x)、loss.backward()）都在 GPU 上进行

什么时候需要 GPU？ 数据量小、模型简单（比如上面的两层 CNN）时，CPU 几秒就能训完，用不用 GPU 区别不大。一旦模型变大（ResNet、Transformer）或数据量上到几万张图，CPU 可能要几小时甚至跑不动，GPU 就是刚需。经验法则：训练时间超过几分钟就考虑上 GPU。纯推理（只跑一次前向）通常 CPU 也够用。

2025 趋势——混合精度与编译加速：现代 GPU（NVIDIA H100 / B200、Apple M4 等）的 float16/bfloat16 算力远超 float32。PyTorch 的 torch.autocast 能自动在合适的地方用半精度计算，显存减半、速度翻倍，几乎不损失精度。torch.compile 则进一步把计算图编译成针对硬件优化的 kernel。两者叠加，训练速度可以提升 3-5 倍。

第 6 步：保存与加载模型

训练好的模型要存下来，否则关掉程序就全丢了。

# —— 方式一：只存权重（推荐）——
# state_dict 是一个字典，键是层名、值是对应的参数 Tensor
torch.save(model.state_dict(), "mnist_cnn.pth")
# 文件很小（几 MB），只包含数值，不包含代码结构

# 加载时必须先创建同样结构的模型，再把权重灌进去
model2 = MNISTClassifier(num_classes=10)        # 先实例化相同结构的空模型
model2.load_state_dict(torch.load("mnist_cnn.pth"))   # 把权重加载进来
model2 = model2.to(device)
model2.eval()   # 加载后记得切换到评估模式

# —— 方式二：存整个模型（不推荐跨环境用）——
torch.save(model, "mnist_cnn_full.pth")
model3 = torch.load("mnist_cnn_full.pth")   # 不需要预先定义类，但依赖 pickle
# 缺点：加载时必须能 import 到原始类的定义，跨 PyTorch 版本或改了类名就会报错

state_dict 是什么？ 它是 PyTorch 模型的”参数字典”——记录了每一层的名字和对应的权重张量，但不包含网络结构的代码。所以推荐的保存流程是：代码里定义模型类（结构）+ 单独存 state_dict（权重），加载时先 new_model = MyModel(...) 建好结构，再 load_state_dict 灌入权重。这种”结构代码 + 权重文件”分离的方式最稳健，也是 HuggingFace 等生态的标准做法。

训练中途做 checkpoint（断点续训）也是同样思路，额外多存一份 optimizer.state_dict()：

# 保存训练检查点（模型权重 + 优化器状态 + 当前 epoch）
checkpoint = {
    "epoch": epoch,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "loss": loss.item(),
}
torch.save(checkpoint, "checkpoint.pth")

# 恢复训练
ckpt = torch.load("checkpoint.pth")
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
# 从 ckpt["epoch"] 继续训练

小结：完整训练流水线

把上面六步串起来，就是一个可以直接复制运行的标准训练脚本骨架：

掌握这个流水线后，你可以把 MNIST 换成 CIFAR-10（记得把输入通道改成 3）、把网络加深、加上 Dropout 防过拟合、换用预训练的 ResNet 做迁移学习——核心套路完全不变。进阶内容见卷积神经网络和迁移学习。

2025 年 PyTorch 前沿：大模型时代的新工具

PyTorch 在 2025 年的进化方向紧紧围绕大模型训练效率。以下是最值得关注的新特性：

torch.compile：从”一行加速”到全栈编译

torch.compile（首次在 PyTorch 2.0 引入）是 PyTorch 从动态图走向”动静态结合”的里程碑。它的原理是：用 TorchDynamo（一个 Python 字节码级别的追踪器）捕获你的前向传播代码，再用 TorchInductor（后端编译器）生成优化的 Triton/CUDA kernel。

到 2025 年的 2.10–2.13 版本，torch.compile 已经大幅成熟：

combo-kernels（2.10）：把多个小 kernel 融合成一个大 kernel，减少 GPU 启动开销。
CUDA Graph 统一 API（2.12）：torch.accelerator.Graph 统一了 CUDA / XPU / 第三方后端的图捕获与重放，减少 kernel launch 延迟。
层次化编译：允许编译大模型的子模块，而非全量编译，降低内存峰值。
Python 3.14 / 3.15 全面支持：包括 free-threaded 构建（无 GIL 模式）。

# 2025 年最佳实践：compile + autocast 混合精度，配合 FSDP2 多卡训练
import torch

model = MyLargeModel().cuda()
model = torch.compile(model, mode="max-autotune")   # 编译加速

with torch.autocast("cuda", dtype=torch.bfloat16):  # 混合精度
    output = model(input)
    loss = criterion(output, target)
loss.backward()
optimizer.step()

FlexAttention：一行代码自定义注意力

注意力机制（Attention）是 Transformer 的核心，但标准实现只支持有限的几种变体（因果、双向等）。2024-2025 年引入的 FlexAttention（灵活注意力，PyTorch 原生的可编程注意力机制）让你用几行代码定义任意注意力模式（如 sliding window、document masking、FlexAR 等），同时自动获得接近 FlashAttention（一种 IO 感知的注意力加速算法）的性能。

到 PyTorch 2.13，FlexAttention 已扩展到 Apple Silicon（MPS 后端，对稀疏注意力模式比 SDPA 快约 12 倍），并在 Hopper / Blackwell GPU 上支持 FlashAttention-4 后端。

大规模训练：FSDP2 与分布式通信

FSDP2（Fully Sharded Data Parallel v2）：PyTorch 全新的分布式训练方案。相比第一代，它基于 per-parameter sharding（逐参数分片），更容易与 torch.compile 配合，调试也更直观。2.13 新增了 reduce-scatter / all-gather 通信重叠（通过独立进程组），进一步提升吞吐。
nn.LinearCrossEntropyLoss（2.13）：把大语言模型最后一层的 Linear + CrossEntropy 融合成一个 op，通过不实体化 logits 张量，将峰值 GPU 显存降低最多 4 倍——对大词表 LLM 训练（如词表 128K+）非常关键。
torchcomms（2.13）：全新的分布式通信后端，提升大规模集群训练的容错性和可调试性。

这些功能主要面向训练数十亿乃至万亿参数大模型的场景。入门阶段不需要用，但了解它们的存在有助于理解 PyTorch 生态的全貌。详见分布式训练。

模型导出与部署：torch.export

PyTorch 2.x 逐步淘汰了 TorchScript（torch.jit），转向基于 torch.export 的新一代导出管线。torch.export 能把动态图模型捕获为一个静态、可序列化的计算图，方便部署到 C++、移动端、ONNX Runtime 等环境。2.12-2.13 还增加了对 Microscaling（MX）量化格式的导出支持。

常用 API 速查

API	用途	示例
`torch.tensor(data)`	从数据创建张量	`torch.tensor([1,2,3])`
`torch.randn(*shape)`	标准正态随机张量	`torch.randn(100,10)`
`torch.zeros/ones(shape)`	全零 / 全一百张量	`torch.zeros(3,4)`
`tensor.to("cuda")`	迁移到 GPU	`x = x.to("cuda")`
`tensor.backward()`	反向传播求梯度	`loss.backward()`
`tensor.detach()`	脱离计算图	`x = x.detach()`
`torch.compile(model)`	编译模型以加速	`model = torch.compile(model)`
`torch.autocast("cuda")`	混合精度上下文	`with torch.autocast("cuda"): ...`
`nn.Linear(in,out)`	全连接层	`nn.Linear(784,128)`
`nn.Sequential(...)`	串联多个层	`nn.Sequential(l1,l2)`
`nn.CrossEntropyLoss()`	多分类损失	`criterion(pred,y)`
`optim.Adam(params,lr)`	Adam 优化器	`optim.Adam(model.parameters())`
`DataLoader(dataset)`	批量加载数据	`DataLoader(ds,batch_size=32)`
`model.train()` / `eval()`	切换训练 / 推理模式	`model.eval()`
`torch.save(obj,path)`	保存模型 / 张量	`torch.save(model.state_dict(),"m.pt")`

实践要点

务必 optimizer.zero_grad()：PyTorch 的梯度默认会累积而非覆盖。如果不清零，每次反向传播的梯度会叠加到上一次的梯度上，导致训练崩溃。这是新手最常踩的坑。
train() 与 eval() 模式有别：Dropout 和 BatchNorm 在训练与推理时行为不同。评估或测试前务必调用 model.eval()，训练时再切回 model.train()。配合 with torch.no_grad() 还能省下推理时的显存。
数据搬到 GPU，模型也要搬：model.to("cuda") 和 X.to("cuda") 缺一不可，否则运算仍跑在 CPU 上。设备不匹配会直接报错。
用 DataLoader 管理批次：直接全量数据训练既慢又爆显存。DataLoader 自动分批、打乱、多进程加载，是标准做法。
混合精度加速训练：torch.autocast 用 float16/bfloat16 计算可提速一倍以上、省一半显存，2025 年已成为训练标配。详见混合精度训练。
用 torch.compile 加速：PyTorch 2.x 的一行编译加速器能自动融合算子、优化 kernel，Transformer 等模型可提速 30%-200%。首次调用有编译开销，适合训练多轮的场景。
从 checkpoint 恢复：保存 model.state_dict()（只存权重）比存整个模型更稳健，跨版本迁移不易出错。详见分布式训练。

典型应用场景

大语言模型训练与微调：GPT、LLaMA、Qwen 等主流 LLM 均基于 PyTorch 生态，HuggingFace Transformers 底层也是它。2025 年的 FSDP2 + FlexAttention + nn.LinearCrossEntropyLoss 组合已成为大规模预训练的标准工具链。详见语言模型演进。
计算机视觉模型：CNN、目标检测、图像分割的模型实现首选 PyTorch。详见卷积神经网络。
Transformer 架构实现：从 BERT 到 ViT，注意力机制的灵活实现得益于 PyTorch 的动态图。详见Transformer 架构。
生成式模型：扩散模型（Stable Diffusion）、GAN、VAE 的研究几乎都以 PyTorch 为基础。详见扩散模型。
迁移学习与微调：预训练模型加载几行代码，在 torchvision.models 或 HuggingFace 上获取。详见迁移学习。

与同类工具对比

特性	PyTorch	TensorFlow 2 / Keras	JAX
计算图	动态（define-by-run）+ `torch.compile`	默认动态（eager）	静态（JIT 编译）
调试体验	极佳，如普通 Python	较好	较难（函数式编程）
API 风格	Pythonic、面向对象	高层 Keras 简洁	函数式
学术占有率	第一（90%+ 论文）	下降中	上升中
工业部署	torch.export / ONNX / TorchServe	TF Serving / TFLite 成熟	较新
大模型生态	HuggingFace 首选	较弱	Google 内部为主
最佳场景	研究 / 大模型 / 灵活实验	生产部署 / 移动端	高性能研究

关键洞察：PyTorch 赢在研究灵活性与生态，TensorFlow 赢在成熟的生产部署工具链。如今学术与大模型领域 PyTorch 已是事实标准，2025 年 torch.compile 与 torch.export 的成熟也让 PyTorch 在生产部署端快速追赶。本站后续框架对比详见TensorFlow/Keras 入门指南。