1. 从一行 import 说起：为什么 PyTorch 值得你花时间

如果你刚接触深度学习，或者从别的框架迁移过来，第一件让你觉得“不太一样”的事情，大概率是这行代码：

import torch

它看起来平平无奇，但背后藏着一整套设计哲学。我见过不少朋友，装完环境跑通一个 MNIST 例子就觉得自己会 PyTorch 了，结果一遇到自定义层、动态 batch、多卡训练就卡壳。问题不在于他们不够聪明，而在于没有理解 PyTorch 到底在解决什么问题。

PyTorch 的核心价值可以用一句话概括： 它让张量计算和自动求导变得像写普通 Python 一样自然 。你定义一个网络，就是在写一个普通的类；你调试一个模型，就是在用 pdb 打断点；你改一个结构，就是在改几行 Python 代码。这种“不隔一层”的体验，是它从学术界一路火到工业界的根本原因。

这篇文章适合谁？如果你满足下面任意一条，接下来的内容会对你有用：

• 刚装好 PyTorch，能跑 demo 但不知道下一步该学什么；
• 从其他框架转过来，想搞清楚 PyTorch 的“脾气”；
• 已经用了一段时间，但训练总出玄学问题，想系统梳理一遍；
• 想了解从数据加载到模型部署的完整链路，而不只是 model.fit() 。

我不会只给你一堆 API 列表，而是把每个环节“为什么这么设计”“实际用的时候哪里会疼”“怎么绕过去”讲清楚。下面从最底层的张量开始，一层层往上走。

2. 张量不只是数组：理解 PyTorch 的数据底座

2.1 张量的三个关键属性：dtype、device、requires_grad

很多人把 PyTorch 张量当成 NumPy 数组的替代品，这个理解只对了一半。张量确实支持索引、切片、广播这些数组操作，但它比 NumPy 多了三个决定性的属性：

• dtype ：数据类型，常见的有 torch.float32 、 torch.float16 、 torch.int64 、 torch.bool 。dtype 选错是新手最常见的报错来源之一，比如把整型张量丢进需要浮点的层里，会直接抛 RuntimeError 。
• device ：张量在 CPU 还是 GPU 上。CPU 张量和 GPU 张量做运算会报错，必须显式 .to(device) 。这个设计是刻意的，避免你无意中把数据在设备间来回拷贝，拖慢训练。
• requires_grad ：这个张量是否需要计算梯度。只有 requires_grad=True 的张量，在参与运算后才会被记录到计算图中，反向传播时才会更新。

我建议你在写任何模型代码之前，先养成一个习惯： 在每个关键节点打印张量的 shape、dtype、device 。这个习惯能帮你省下大量调试时间。比如下面这段：

x = torch.randn(4, 3, 224, 224)
print(x.shape, x.dtype, x.device)
# torch.Size([4, 3, 224, 224]) torch.float32 cpu

看起来简单，但当你处理一个真实数据集时，图像可能是 uint8 ，标签可能是 int64 ，模型权重是 float32 ，三者混在一起运算就会出问题。提前检查，比事后猜错要高效得多。

2.2 广播机制：方便，但也是 bug 温床

PyTorch 的广播规则和 NumPy 一致：从最后一个维度开始对齐，维度大小为 1 或缺失的维度可以自动扩展。这让你能写出很简洁的代码，比如：

a = torch.randn(4, 1, 5)
b = torch.randn(3, 5)
c = a + b  # 结果是 (4, 3, 5)

但广播也是很多隐蔽 bug 的来源。我踩过最典型的一个坑是： 标签张量和预测张量形状不匹配，广播后 loss 看起来在下降，但模型根本没学到东西 。比如预测是 (batch, 1) ，标签是 (batch,) ，两者相减会广播成 (batch, batch) ，算出来的 loss 完全错误，但程序不会报错。

提示：在计算 loss 之前，养成检查 pred.shape 和 target.shape 的习惯。如果两者维度不一致，先想清楚是应该 squeeze 、 unsqueeze 还是 view ，而不是依赖广播“自动修复”。

2.3 view、reshape、permute 的区别与选择

这三个操作经常被混用，但它们的语义不同：

操作	是否要求内存连续	是否拷贝数据	典型用途
view	是	否	改变形状，要求原张量连续
reshape	否	视情况	改变形状，不连续时会拷贝
permute	否	否	交换维度顺序，返回非连续视图
transpose	否	否	交换两个维度

实际使用中， permute 之后直接 view 会报错，因为 permute 返回的张量在内存中不连续。正确做法是 permute 之后接 .contiguous() 再 view ，或者直接用 reshape 。这个细节在处理图像通道变换（如 (B, H, W, C) 转 (B, C, H, W) ）时经常遇到。

3. 自动求导：PyTorch 最核心的那层“魔法”

3.1 计算图是怎么建起来的

PyTorch 的自动求导基于动态计算图。每次你对 requires_grad=True 的张量做运算，PyTorch 就会在后台记录这个操作，形成一个有向无环图。当你调用 .backward() 时，它沿着这个图反向走一遍，用链式法则算出每个叶子节点的梯度。

这个过程的关键在于： 图是在前向传播时动态构建的，每次迭代都会重新建图 。这和静态图框架（先定义图再喂数据）有本质区别。动态图的好处是调试方便，你可以像写普通 Python 一样加 if 、 for 、 print ；代价是每次迭代都有建图开销，不过在现代硬件上这个开销通常可以接受。

一个常见的误解是： loss.backward() 之后梯度会一直累积。实际上 PyTorch 的梯度是 累加 的，不是覆盖的。所以标准训练循环里必须写：

optimizer.zero_grad()
loss.backward()
optimizer.step()

如果你忘了 zero_grad() ，梯度会越加越大，训练直接发散。这个坑几乎每个新手都会踩一次。

3.2 什么时候该用 torch.no_grad()

推理阶段不需要梯度，这时候用 torch.no_grad() 包起来，能显著减少显存占用和计算量：

model.eval()
with torch.no_grad():
    for x, y in test_loader:
        pred = model(x)
        # 计算指标

注意 model.eval() 和 torch.no_grad() 是两件事：前者切换 BatchNorm 和 Dropout 的行为，后者关闭梯度记录。推理时两个都要用，缺一不可。我见过有人只写了 no_grad 但忘了 eval() ，结果 BatchNorm 用训练模式的统计量，指标怎么都对不上。

3.3 梯度裁剪与梯度检查

训练 RNN、Transformer 这类深层网络时，梯度爆炸是常见问题。PyTorch 提供了 torch.nn.utils.clip_grad_norm_ 和 clip_grad_value_ 两个工具：

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

max_norm 设多少？经验值是 0.5 到 5.0 之间，具体要看任务。我一般从 1.0 开始试，如果训练不稳定就往下调。另外，如果你想检查梯度是否正常，可以在 backward() 之后遍历参数打印梯度范数：

for name, param in model.named_parameters():
    if param.grad is not None:
        print(name, param.grad.norm().item())

如果某个层的梯度范数是 0 或者 NaN，说明那里出了问题，可能是初始化不当，也可能是激活函数饱和。

4. 模型构建：从 nn.Module 到自定义层

4.1 nn.Module 的约定与陷阱

PyTorch 里所有模型都继承 nn.Module 。它有两个必须遵守的约定：

1. 子模块必须赋值给 self 的属性 ，这样 parameters() 才能找到它们；
2. forward 方法定义前向逻辑 ，不要在里面做有状态的操作。

看一个反例：

class BadModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(10, 10) for _ in range(3)]  # 错误：列表不是 Module

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

这样写 model.parameters() 返回空，优化器拿不到参数，训练完全不动。正确做法是用 nn.ModuleList 或 nn.Sequential ：

self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(3)])

这个坑非常隐蔽，因为代码能跑通，loss 也能算出来，但参数永远不更新。我建议你在定义完模型后，打印一下 sum(p.numel() for p in model.parameters()) ，确认参数量符合预期。

4.2 自定义层：什么时候需要，怎么写

内置层不够用时，就需要自定义。比如你想实现一个带可学习温度的 Softmax：

class TemperatureSoftmax(nn.Module):
    def __init__(self, dim=-1, init_temp=1.0):
        super().__init__()
        self.dim = dim
        self.temperature = nn.Parameter(torch.tensor(init_temp))

    def forward(self, x):
        return torch.softmax(x / self.temperature, dim=self.dim)

这里 nn.Parameter 把普通张量注册成可学习参数，优化器会自动包含它。注意温度不能为负，实际使用中可以加个 softplus 或者 clamp 保证正值。

自定义层的价值在于把领域知识编码进模型结构。比如做时序任务时，你可能需要一个只关注局部窗口的注意力层；做图任务时，需要一个基于邻接矩阵的消息传递层。这些都不是内置层能直接覆盖的。

4.3 权重初始化：默认值够用吗

PyTorch 的默认初始化是 Kaiming 均匀分布（对 Linear 和 Conv），大多数情况下够用。但如果你用了自定义激活函数，或者网络特别深，可能需要手动初始化：

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model.apply(init_weights)

apply 会递归遍历所有子模块。注意初始化要在模型创建之后、训练之前做，而且不要覆盖预训练权重。

5. 数据管道：Dataset 和 DataLoader 的正确打开方式

5.1 Dataset 的三种写法

PyTorch 提供两种 Dataset： map-style 和 iterable-style 。绝大多数场景用 map-style，需要实现 __len__ 和 __getitem__ ：

class MyDataset(Dataset):
    def __init__(self, paths, labels, transform=None):
        self.paths = paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        img = load_image(self.paths[idx])
        if self.transform:
            img = self.transform(img)
        return img, self.labels[idx]

关键点是 __getitem__ 只负责取 一个 样本，不要在里面做 batch 操作。batch 由 DataLoader 负责。

如果数据量特别大，放不进内存，可以用 IterableDataset ，它实现 __iter__ 而不是 __getitem__ ，适合流式读取。但要注意多进程下每个 worker 会独立迭代，需要手动分片。

5.2 DataLoader 的参数怎么调

DataLoader 有几个参数直接影响训练效率和正确性：

• batch_size ：受显存限制，一般从 32 或 64 开始试。太小导致梯度噪声大，太大导致泛化变差。
• shuffle ：训练集必须 True，验证集和测试集必须 False。忘了设 False 会导致评估指标不稳定。
• num_workers ：数据加载的并行进程数。设成 CPU 核心数通常不错，但 Windows 上多进程有额外开销，建议从 0 开始试。
• pin_memory ：GPU 训练时设为 True，能加速 CPU 到 GPU 的拷贝。
• drop_last ：训练时设为 True 可以避免最后一个不完整 batch 影响 BatchNorm。

一个我常用的配置：

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True
)

persistent_workers=True 在 epoch 之间保持 worker 存活，避免反复创建销毁的开销，数据量大时效果明显。

5.3 数据增强放在哪里

数据增强应该放在 Dataset 的 __getitem__ 里，用 torchvision.transforms 或者自定义函数。注意增强只对训练集做，验证集只做必要的 resize 和归一化。

一个容易忽略的点是： 归一化的均值和标准差要用训练集的统计量 ，不能各自算各自的。否则验证集分布和训练集不一致，指标会虚高。

6. 训练循环：那些教程不会告诉你的细节

6.1 标准训练循环的骨架

for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        pred = model(x)
        loss = criterion(pred, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

    model.eval()
    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            pred = model(x)
            # 累计指标

这个骨架看起来简单，但每个位置都有讲究。 zero_grad 必须在 backward 之前， clip_grad 必须在 step 之前， train() 和 eval() 必须成对出现。

6.2 学习率调度：什么时候降，降多少

固定学习率很少是最优解。PyTorch 在 torch.optim.lr_scheduler 里提供了多种调度器：

调度器	行为	适用场景
StepLR	每 N 个 epoch 乘以 gamma	传统 CNN 训练
CosineAnnealingLR	余弦曲线下降	Transformer、现代网络
ReduceLROnPlateau	指标不降时降	不确定何时该降
OneCycleLR	先升后降	快速收敛

我个人的偏好是： 不确定就用 CosineAnnealingLR，配合 warmup 。warmup 在前几个 epoch 把学习率从很小线性升到目标值，能避免训练初期的不稳定。

注意调度器的 step() 调用位置：大多数调度器每个 epoch 调一次，但 OneCycleLR 是每个 batch 调一次。调错了学习率曲线就完全不对。

6.3 混合精度训练：省显存提速的实用技巧

torch.cuda.amp 让混合精度训练变得很简单：

scaler = torch.cuda.amp.GradScaler()

for x, y in loader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        pred = model(x)
        loss = criterion(pred, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

关键点是 scale 和 step 的配合： GradScaler 会自动调整损失缩放因子，防止梯度下溢。用混合精度后显存通常能省 30% 到 50%，速度提升 20% 到 40%，代价是数值精度略降，大多数任务不受影响。

但要注意： 有些操作在 fp16 下不稳定 ，比如 softmax 的指数运算、大数相减。如果遇到 NaN，先关掉 amp 确认是不是精度问题。

7. 显存不够用？先别急着换卡

7.1 显存都去哪了

训练时显存主要被四部分占用：

1. 模型参数 ：参数量乘以 4 字节（fp32）；
2. 梯度 ：和参数量相同；
3. 优化器状态 ：Adam 会存两份动量，是参数量的两倍；
4. 激活值 ：前向传播的中间结果，batch 越大占用越多。

一个 1000 万参数的模型，用 Adam 训练，光参数、梯度、优化器状态就要 1000万 × 4 × 4 = 160MB，激活值往往是大头。

7.2 梯度累积：小显存跑大 batch

如果 batch_size 受显存限制，可以用梯度累积模拟大 batch：

accum_steps = 4
for i, (x, y) in enumerate(loader):
    loss = model(x, y) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

注意 loss 要除以累积步数，否则梯度会放大。另外 BatchNorm 的统计量是按实际 batch 算的，梯度累积不能完全等价于大 batch，这点要心里有数。

7.3 检查点与梯度检查点

torch.utils.checkpoint 可以用计算换显存：前向时不保存中间激活，反向时重新算一遍。代价是训练速度慢 20% 到 30%，但显存能省很多。适合深层网络。

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    x = checkpoint(self.block1, x)
    x = checkpoint(self.block2, x)
    return x

注意 checkpoint 的输入必须 requires_grad=True ，否则不生效。

8. 从训练到部署：模型保存与加载的坑

8.1 state_dict 还是整个模型

PyTorch 有两种保存方式：

# 方式一：只存参数
torch.save(model.state_dict(), "model.pth")

# 方式二：存整个模型
torch.save(model, "model.pth")

推荐方式一 。存整个模型依赖类的定义路径，换环境或改代码后可能加载失败。 state_dict 只存张量，更稳定。

加载时：

model = MyModel()
model.load_state_dict(torch.load("model.pth", map_location="cpu"))
model.eval()

map_location 很重要：在 GPU 上存的模型，加载到 CPU 环境时必须指定，否则会报错。

8.2 推理时的常见问题

导出模型给其他环境用时， torch.jit.trace 或 torch.jit.script 能把模型转成 TorchScript：

model.eval()
example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)
traced.save("model.pt")

trace 只记录实际执行路径，如果模型有 if 分支，可能漏掉某些路径。 script 会解析源码，支持控制流，但对动态特性支持有限。选择哪个取决于模型复杂度。

另外，推理前一定要 model.eval() ，并且用 torch.no_grad() 包住，否则显存会持续增长。

9. 我踩过的几个真实坑

第一个坑是 DataLoader 的 num_workers 在 Windows 上导致死锁 。当时设了 num_workers=4 ，程序卡在第一个 batch 不动。后来改成 0 就正常了。Windows 的多进程启动方式和 Linux 不同，如果非要用多 worker，需要把训练代码放在 if __name__ == "__main__": 里面。

第二个坑是 BatchNorm 在 batch_size=1 时出错 。BN 需要至少两个样本才能算方差，batch 为 1 时要么报错，要么用 track_running_stats=False 。小 batch 场景建议用 GroupNorm 或 LayerNorm 替代。

第三个坑是 学习率设太大导致 loss 变 NaN 。排查了半天以为是数据问题，最后发现是初始学习率 0.1 太大。现在我的习惯是先用小学习率跑几个 batch，确认 loss 正常下降再放大。

第四个坑是 忘记设随机种子导致结果不可复现 。完整设置是：

import random
import numpy as np
import torch

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

注意 deterministic=True 会降低速度，只在需要严格复现时开。

10. 接下来可以往哪走

PyTorch 的生态很大，训练只是其中一环。如果你已经把训练跑通，下一步可以看看这几个方向：

• 分布式训练 ： DistributedDataParallel 是多卡训练的标准方案，比 DataParallel 更高效；
• 性能分析 ： torch.profiler 能定位瓶颈在数据加载、前向还是反向；
• 模型压缩 ：量化、剪枝、蒸馏，把大模型塞进小设备；
• 部署工具链 ：TorchScript、ONNX 导出、移动端推理。

我的建议是不要贪多，先把单卡训练这条链路走扎实。能稳定复现一个任务、能定位训练中的问题、能根据指标调整超参，这三点做到了，再往深处走会顺很多。PyTorch 的文档和源码质量都很高，遇到不确定的地方，直接看源码往往比搜博客更快。