PyTorch 学习

本文最后更新于 2026-08-14

PyTorch 学习的关键不是记住所有张量函数，而是建立一条可重复使用的工程主线：准备数据、构建模型、计算损失、反向传播、更新参数、验证模型、保存权重。

本文专注 PyTorch 的实现方式。激活函数、初始化、优化器和正则化背后的原理见《深度学习总结》；特征工程和模型评估见《机器学习总结》。

学习路线

阶段 需要掌握的内容 学完后的目标
张量 shape、dtype、device、索引和广播 能判断每一步的数据形状
自动微分 requires_grad、backward()、梯度累积 能解释参数梯度从哪里来
数据管线 Dataset、DataLoader、batch 能稳定地向模型提供数据
模型 nn.Module、forward()、参数 能搭建并检查网络
训练与评估 train()、eval()、优化器、无梯度推理 能写出不混用数据的训练循环
持久化 state_dict、checkpoint、设备映射 能恢复训练并部署推理

环境与设备

安装和版本检查

CPU、CUDA 和不同操作系统对应的安装命令可能不同，应先使用 PyTorch 安装选择器 生成命令。安装后检查环境：

1
2
3
4
5
import torch

print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("MPS available:", torch.backends.mps.is_available())

统一选择设备

1
2
3
4
5
6
7
8
9
10
def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


device = get_device()
print("device:", device)

模型和参与同一次运算的张量必须位于同一设备：

1
2
3
model = model.to(device)
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)

不要在循环中频繁地把张量来回移动到 CPU 和 GPU，这会增加同步与传输开销。

可复现性

1
2
3
4
5
6
7
8
9
10
11
12
import random

import numpy as np
import torch

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

固定随机种子有助于复现实验，但不同硬件、算子和 PyTorch 版本仍可能产生细微差异。需要严格确定性时，还要检查所用算子是否支持 deterministic 模式。

张量基础

创建、dtype 与 device

1
2
3
4
5
6
7
8
9
import torch

x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
zeros = torch.zeros((2, 3))
ones = torch.ones_like(zeros)
random_values = torch.randn((2, 3))
sequence = torch.arange(0, 10, 2)

print(x.shape, x.dtype, x.device)

模型参数通常使用浮点类型；分类标签传给 CrossEntropyLoss 时通常使用 torch.long。许多难以理解的报错，本质上来自 shape、dtype 或 device 不一致。

形状操作

操作 作用 关键区别
reshape() 改变视图形状 必要时会复制数据，通用性更好
view() 按新形状查看同一块数据 通常要求内存连续
transpose() 交换两个维度 常导致张量不连续
permute() 按指定顺序重排多个维度 常用于通道顺序转换
unsqueeze() 增加大小为 1 的维度 常用于补 batch 或 channel 维
squeeze() 删除大小为 1 的维度 最好指定维度，避免误删 batch 维 
1
2
3
4
5
x = torch.arange(24).reshape(2, 3, 4)

print(x.shape)                     # [2, 3, 4]
print(x.permute(0, 2, 1).shape)    # [2, 4, 3]
print(x.flatten(start_dim=1).shape) # [2, 12]

拼接、堆叠与广播

  • torch.cat() 沿已有维度拼接，除拼接维度外形状必须一致。
  • torch.stack() 创建新维度，所有输入张量形状必须完全一致。
  • 广播允许不同形状参与运算，但维度需要从末尾开始兼容。
1
2
3
4
5
6
a = torch.ones((2, 3))
b = torch.zeros((2, 3))

print(torch.cat([a, b], dim=0).shape)   # [4, 3]
print(torch.stack([a, b], dim=0).shape) # [2, 2, 3]
print((a + torch.tensor([1, 2, 3])).shape)  # [2, 3]

Tensor、NumPy 与 Python 标量

CPU Tensor 和 NumPy 数组可能共享内存：

1
2
3
4
5
6
import numpy as np

array = np.array([1, 2, 3], dtype=np.float32)
tensor = torch.from_numpy(array)
back_to_numpy = tensor.numpy()
scalar = tensor[0].item()

需要从计算图中取出张量用于日志或 NumPy 处理时，常用：

1
values = tensor.detach().cpu().numpy()

不要使用 .data 绕过 Autograd，它可能让梯度计算悄悄出错。

Autograd 自动微分

PyTorch 在前向计算时动态记录计算图，调用 backward() 后根据链式法则计算梯度。

1
2
3
4
5
6
7
8
9
10
11
12
x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

y = w * x + b
loss = (y - 10) ** 2
loss.backward()

print("loss:", loss.item())
print("dx:", x.grad)
print("dw:", w.grad)
print("db:", b.grad)

需要记住四点：

  1. 只有浮点或复数张量才能计算梯度。
  2. 叶子张量的梯度保存在 .grad 中。
  3. 多次 backward() 会累积梯度，训练时需要显式清零。
  4. detach() 返回与原张量共享数据但不继续跟踪当前计算图的张量。

推理阶段可使用 torch.inference_mode()，减少梯度记录和额外开销：

1
2
3
model.eval()
with torch.inference_mode():
    predictions = model(features)

Dataset 与 DataLoader

Dataset 表示“如何获取一个样本”，DataLoader 负责 batching、shuffle 和并行加载。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from torch.utils.data import DataLoader, TensorDataset

x = torch.randn(100, 5)
y = torch.randn(100, 1)

dataset = TensorDataset(x, y)
loader = DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,
)

for x_batch, y_batch in loader:
    print(x_batch.shape, y_batch.shape)
    break

只对训练集启用 shuffle=True。验证集和测试集通常不需要打乱；更重要的是它们必须由独立数据创建，不能误用训练集。

使用 nn.Module 构建模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from torch import nn


class RegressionMLP(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(in_features, 64),
            nn.ReLU(),
            nn.Dropout(p=0.1),
            nn.Linear(64, 1),
        )

    def forward(self, x):
        return self.network(x)

在 __init__() 中注册层，在 forward() 中定义数据流。直接调用 model(x)，不要手动调用 model.forward(x)，因为 nn.Module 的调用过程还负责 hooks 等框架逻辑。

检查模型时重点关注：

1
2
3
4
model = RegressionMLP(in_features=5)
print(model)
print("parameters:", sum(p.numel() for p in model.parameters()))
print("trainable:", sum(p.numel() for p in model.parameters() if p.requires_grad))

标准训练与评估循环

单轮训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def train_one_epoch(model, dataloader, loss_fn, optimizer, device):
    model.train()
    total_loss = 0.0
    total_samples = 0

    for features, targets in dataloader:
        features = features.to(device)
        targets = targets.to(device)

        optimizer.zero_grad(set_to_none=True)
        predictions = model(features)
        loss = loss_fn(predictions, targets)
        loss.backward()
        optimizer.step()

        batch_size = features.size(0)
        total_loss += loss.item() * batch_size
        total_samples += batch_size

    return total_loss / total_samples

zero_grad()、backward()、step() 是每个训练 step 的核心。这里按样本数加权汇总损失，避免最后一个不完整 batch 造成平均值偏差。

验证循环

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
@torch.inference_mode()
def evaluate(model, dataloader, loss_fn, device):
    model.eval()
    total_loss = 0.0
    total_samples = 0

    for features, targets in dataloader:
        features = features.to(device)
        targets = targets.to(device)
        predictions = model(features)
        loss = loss_fn(predictions, targets)

        batch_size = features.size(0)
        total_loss += loss.item() * batch_size
        total_samples += batch_size

    return total_loss / total_samples

model.eval() 会切换 Dropout 和 BatchNorm 等层的行为；inference_mode() 负责关闭梯度记录。两者解决的问题不同，都不能省略。

完整回归案例

下面使用合成数据演示完整流程。训练集和验证集严格分开，最后同时返回两组 loss。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(42)
device = get_device()

# y = 2.5 * x + 5.2 + noise
x = torch.randn(200, 1)
y = 2.5 * x + 5.2 + torch.randn(200, 1) * 0.3

indices = torch.randperm(len(x))
train_indices = indices[:160]
val_indices = indices[160:]

train_dataset = TensorDataset(x[train_indices], y[train_indices])
val_dataset = TensorDataset(x[val_indices], y[val_indices])
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

model = nn.Linear(1, 1).to(device)
loss_fn = nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.05)

train_losses = []
val_losses = []

for epoch in range(1, 101):
    train_loss = train_one_epoch(
        model, train_loader, loss_fn, optimizer, device
    )
    val_loss = evaluate(model, val_loader, loss_fn, device)
    train_losses.append(train_loss)
    val_losses.append(val_loss)

    if epoch == 1 or epoch % 20 == 0:
        print(
            f"epoch={epoch:03d} "
            f"train_loss={train_loss:.4f} "
            f"val_loss={val_loss:.4f}"
        )

print("weight:", model.weight.item())
print("bias:", model.bias.item())

一个可靠的训练案例至少应该满足：数据划分正确、训练和评估模式正确、验证阶段不计算梯度、分别记录训练与验证指标。

参数初始化

PyTorch 的标准层已经有默认初始化。确实需要自定义时，可以通过 apply() 统一处理指定层：

1
2
3
4
5
6
7
8
def initialize_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_uniform_(module.weight, nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)


model.apply(initialize_weights)

初始化方法要与激活函数和网络结构匹配，不应仅因为“常用”就覆盖所有层。Xavier、Kaiming 以及对称性问题见《深度学习总结》。

保存和恢复模型

只保存推理权重

1
2
3
4
5
6
7
8
9
10
torch.save(model.state_dict(), "model_weights.pth")

restored_model = nn.Linear(1, 1).to(device)
state_dict = torch.load(
    "model_weights.pth",
    map_location=device,
    weights_only=True,
)
restored_model.load_state_dict(state_dict)
restored_model.eval()

保存 state_dict 比直接序列化完整模型更灵活，也减少了对原 Python 类路径的依赖。

保存可恢复训练的 checkpoint

1
2
3
4
5
6
7
8
9
torch.save(
    {
        "epoch": 100,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "val_loss": val_losses[-1],
    },
    "checkpoint.pth",
)

恢复训练时还应恢复优化器状态、当前 epoch、学习率调度器以及随机状态。若只恢复模型参数而重新创建优化器，Momentum 或 Adam 的历史统计会丢失。

常见错误清单

问题 结果 正确做法
验证 DataLoader 使用训练集 验证结果失真 分别创建 train_dataset 和 val_dataset
验证循环仍遍历训练 loader 实际没有验证 函数参数只接收当前要评估的 loader
返回两次 val_losses 调用方拿不到训练曲线 明确返回 train_losses, val_losses
忘记 optimizer.zero_grad() 梯度跨 batch 累积 每个训练 step 显式清零
忘记 model.eval() Dropout、BatchNorm 推理不稳定 验证和推理前切换模式
忘记关闭验证梯度 额外占用内存和计算 使用 torch.inference_mode()
CPU 与 GPU 张量混用 运行时报 device 错误 模型、输入和标签放到同一设备
无条件 squeeze() batch size 为 1 时误删 batch 维 指定维度，例如 squeeze(-1)
使用 .data 修改张量 破坏 Autograd 正确性 使用 detach() 或 no_grad()

参考资料