最近好几个做模型部署和算法岗的朋友问我，想系统地学深度学习框架，到底应该从哪个入手。我每次的回答都很一致：如果只能选一个，那就 PyTorch。这个结论不是我拍脑袋，而是这几年实际做项目、复现论文、搬模型踩坑踩出来的。PyTorch 的生态、社区、设计思路，已经事实上成了深度学习研究和工业落地的通用语言。今天我就把自己从张量操作到完整训练流程的经验整理出来，尽量讲清楚每个环节背后的道理，而不是只给一堆能跑的代码。
这篇内容适合刚接触深度学习、想用 PyTorch 跑通第一个模型的人，也适合用过其他框架但想转到 PyTorch 生态的开发者。我会从核心设计思路讲起，一步步到数据加载、模型构建、训练循环、断点续训和常见 bug 排查，全程用可复现的代码片段和实际参数说明，争取让你看完能直接上手。
1. 为什么 PyTorch 会成为首选：核心设计思路拆解
1.1 动态计算图：PyTorch 的灵魂所在
深度学习框架最核心的抽象就是计算图。PyTorch 默认采用动态计算图（Define-by-Run），意思是计算图在每次前向传播时边运行边构建，图的结构完全由代码执行路径决定。这和静态图（Define-and-Run）有本质区别，静态图要求你先定义好完整的计算结构再执行，一旦结构变化就要重新编译。
动态图带来的直接好处是调试极其方便。我可以在任意一行代码处打断点，打印中间张量的 shape 和数值，就像调试普通 Python 程序一样。这在研究新模型结构时是救命级的体验，特别是在写带条件分支的模型时，静态图需要把分支全部建模进去，而 PyTorch 里就是普通的 if-else，代码长什么样，计算流就是什么样。
另一个容易被低估的点是：动态图让模型结构和控制流解耦。你写一个包含循环和条件判断的模型，不用额外封装控制流算子，PyTorch 的 autograd 机制会自动跟踪每个张量的操作历史。这意味着模型结构可以跟业务逻辑一样灵活，不再受“先定义后执行”的限制。
1.2 Tensor 统一了计算与存储
PyTorch 里的 Tensor 不仅仅是存储数据的容器，它统一了 CPU、GPU、模型参数和梯度这四种角色。我第一次用 PyTorch 时最强烈的感受是：不需要区分“计算用的数组”和“工程里的大数组”，Tensor 就是一切。CPU 上的 Tensor 可以通过.cuda()或.to(device)一行迁移到 GPU，模型参数本身就是 Tensor，梯度也是 Tensor，这极大地降低了显式数据搬运的成本。
有个细节值得特别说明：Tensor 的设计借鉴了 NumPy，但做了两个关键增强。第一是设备无关性，同一个 API 在 CPU 和 CUDA 上表现一致；第二是自动微分，每个 Tensor 可以通过requires_grad=True开启梯度追踪，默认情况下requires_grad=False，避免不必要的显存和计算开销。我见过很多新手的写法是创建 Tensor 时习惯性全部打开requires_grad，结果发现显存不够用，就是没理解默认值的设计逻辑：只有训练参数才需要梯度，数据输入不需要。
1.3 Pythonic 的学习曲线
PyTorch 的设计哲学是“尽可能让 Python 代码就是模型代码”。你不需要学习一套专门的配置文件语法或图描述语言，模型就是一个torch.nn.Module的子类，前向传播就是重写forward方法。这种风格大幅降低了学习成本，也让代码审查和协作变得容易。
实际经验告诉我，PyTorch 项目里的代码可读性普遍比静态图框架的高。因为你不需要去来回对照“配置文件里的结构”和“实际运行时的结构”，类属性和网络层一一对应，注册子模块用self.xxx = nn.Linear(...)即可，Module会自动追踪所有子模块和参数。这背后是nn.Module的__setattr__魔法方法在起作用：当你把nn.Linear实例赋给self.fc1时，会自动注册到_modules字典，方便后续参数初始化和模型保存。
2. 核心细节解析：张量操作、自动求导与模块体系
2.1 张量创建与基础操作的“肌肉记忆”
我建议所有初学者把下面这段代码当作最基本的肌肉记忆训练，先理解每个函数的适用场景，再谈跑模型。
import torch # 从列表或 NumPy 数组创建 t1 = torch.tensor([1, 2, 3]) t2 = torch.from_numpy(np.array([1, 2, 3], dtype=np.float32)) # 常用初始化 t3 = torch.zeros(2, 3) # 全 0 t4 = torch.ones(2, 3) # 全 1 t5 = torch.randn(2, 3) # 标准正态分布 t6 = torch.randint(0, 10, (2, 3)) # 均匀整数 # 形状操作三件套 t7 = t5.view(3, 2) # 只改变视图，要求内存连续 t8 = t5.reshape(3, 2) # 不要求连续，必要时复制 t9 = t5.permute(1, 0) # 维度交换，常用在图片 HWC 与 CHW 转换 # 拼接与拆分 t10 = torch.cat([t5, t6.float()], dim=0) # 拼接 t11, t12 = torch.chunk(t10, 2, dim=0) # 均分 t13 = torch.split(t10, 1, dim=0) # 按块数拆分这里有三个高频错误我在实际项目里见过好几次。一是view和reshape混用：view要求底层内存连续，而经过transpose或permute后的张量不满足这个条件，直接view会报错或得到错误结果；reshape内部会处理这种情况，代价是可能发生数据拷贝。二是cat前后的 shape 对齐问题：除了拼接所在维度，其他维度必须完全一致，否则报错信息虽然容易看懂，但在循环里反复出错会影响调试效率。三是tensor.item()和tensor.numpy()的使用场景没分清：item()用于取单元素标量，numpy()用于转 NumPy 数组，但带梯度的张量需要先detach()再转。
2.2 自动求导的内功心法：autograd 到底在追踪什么
PyTorch 的自动微分是理解整个框架的关键。它的核心机制是记录“反向传播图”：每个需要梯度的 Tensor 通过requires_grad=True进入追踪，当你在这个张量上执行运算时，PyTorch 不仅算出结果，还会把运算方保存起来，形成一棵带有梯度函数的计算图，反向传播时从根节点往回遍历，用链式法则求每个叶子节点的梯度。
我用一个非常小的例子来说明：
# 假设我们要算 y = x^2 + 3x，求 x=2 处的 dy/dx x = torch.tensor(2.0, requires_grad=True) y = x ** 2 + 3 * x y.backward() print(x.grad) # 期望输出 7.0（2*2 + 3）这段代码背后发生的事情是：y是x的函数，backward()从y出发，沿着计算图反向传播，把梯度累积到x.grad里。注意x.grad是累积的，如果多次调用backward()，梯度会累加而不是覆盖，这在某些循环训练场景下会引发问题，你需要手动x.grad.zero_()。
我最想提醒的一点是：默认情况下backward()的梯度是累加的，所以训练循环里每个 step 必须调用optimizer.zero_grad()，否则梯度会在多次迭代间越积越大，模型参数更新方向完全混乱。如果是自定义循环里忘记清零，loss 曲线会呈现奇怪的锯齿状或直接发散。
2.3 nn.Module 的模块化管理
PyTorch 用nn.Module作为一切模型的基类。它完成的三件关键事情是：子模块注册、参数管理和状态切换。
import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(784, 256) self.relu = nn.ReLU() self.fc2 = nn.Linear(256, 10) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x net = MyNet() print(net) # 会递归打印所有子模块 print([name for name, p in net.named_parameters()]) # 查看所有参数名这里的nn.Linear会自动创建weight和bias两个参数，并注册为模型的子参数。PyTorch 的nn.Sequential适合层数较少的直线结构，而更复杂的结构需要手动管理子模块。我个人的建议是：哪怕只有三层的线性堆叠，也尽量写成类形式，因为后续你需要添加dropout、BatchNorm或者条件分支时，类形式改起来比Sequential灵活得多。
Module的train()/eval()状态切换是另一个容易踩坑的点。BatchNorm和Dropout在训练和推理时的行为完全不同：train()模式下Dropout会随机置零，BatchNorm会使用当前 batch 的统计量；eval()模式下Dropout失效，BatchNorm会使用训练阶段累积的全局统计量。如果推理前忘了调用model.eval()，结果会飘忽不定，而且很难排查。
3. 实操过程：从数据流水线到完整训练循环
3.1 数据加载与预处理：Dataset 和 DataLoader 的正确用法
真实项目里数据加载的坑往往比模型本身还多。PyTorch 的Dataset和DataLoader分工明确：Dataset负责定义“如何读取一条样本”，DataLoader负责“如何组织批量、打乱和并行”。
import torch from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, images, labels): self.images = images # 假设 shape: [N, C, H, W] self.labels = labels def __len__(self): return len(self.images) def __getitem__(self, idx): # 实际项目中这里通常会做图像读取、数据增强、归一化 image = self.images[idx].float() / 255.0 # 归一化到 [0,1] label = self.labels[idx] return image, label train_ds = MyDataset(train_imgs, train_labels) train_loader = DataLoader( train_ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, )DataLoader的几个参数值得逐个解释。shuffle=True在训练时必须开启，让每个 epoch 内样本顺序打乱，防止模型学到样本间的顺序依赖；num_workers>0使用多进程预读取数据，但进程数不是越多越好，实际经验是 4 到 8 在大多数机器上收益最大，超过 CPU 核数反而会因进程切换开销变慢；pin_memory=True在 GPU 训练时建议开启，它让数据的页锁定内存直接传输到 GPU，减少一次拷贝，这个提速在数据量大时非常明显；drop_last=True丢弃最后不足一个 batch 的数据，除特殊情况外建议开启，不然最后一个 batch 的统计量会影响BatchNorm。
关于数据增强，我特别想说一点：在线增强（训练时每轮随机变换）和离线增强（预处理时翻倍数据集）效果差异很大。在线增强的样本多样性更高，因为同一个样本每个 epoch 看到的都可能不同，而且不需要额外存放大规模增强后的数据。对图像任务，常见的增强包括随机裁剪、随机翻转、颜色抖动、旋转，用torchvision.transforms可以很方便地串联。
3.2 构建一个真实可训练的卷积网络
光看 Linear 结构还不够，我顺手给一个经典的 CNN 结构，这是纯新手从分类任务起步最常见的骨架：
import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes=10): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(128 * 4 * 4, 256), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(256, num_classes), ) def forward(self, x): return self.classifier(self.features(x))这里有几处经验选择。padding=1保证 3x3 卷积不改变特征图尺寸，这样整个features部分只需关注池化层对尺寸的缩减。输入如果是32x32，经过三次MaxPool2d(2)后，到 classifier 之前特征图是4x4。很多新手在算Linear的输入维度时算错，我的建议是别急着手动算，先构造一个假张量model.features(torch.zeros(1,3,32,32)).shape打印出来再填Linear的输入维度，这是在深度学习里非常实用的调试手段。
inplace=True这个参数也值得展开说。它表示在输入张量原地执行ReLU而不申请新内存，能省一些显存，代价是如果你后续需要复用该张量的原始值，就会出问题。在forward里用inplace=True一般安全，但如果后面接nn.Sequential且某个中间张量还需要被后续模块引用，就不要开。我见过有人在ReLU(inplace=True)之后又对输入做了别的操作，导致数值变化，本质是原地修改破坏了原始引用。
3.3 训练循环的写法与参数解读
PyTorch 没有内置训练循环，必须自己写，但这反而是件好事，因为你能清楚看到每一步在做什么。一个标准的训练循环骨架如下：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = SimpleCNN(num_classes=10).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20) num_epochs = 20 for epoch in range(num_epochs): model.train() running_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() # 可选：梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() running_loss += loss.item() * images.size(0) scheduler.step() epoch_loss = running_loss / len(train_loader.dataset) print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}") # 验证 model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f"Val Acc: {100 * correct / total:.2f}%")训练循环里的每一步都不能省，但有几个细节我想单独强调。
第一个是loss.item()的用法。如果直接打印loss这个 Tensor，会把整个计算图的信息都打出来，既慢又占内存；.item()只取标量值，兼顾性能和可读性。loss.item() * images.size(0)是计算整个 epoch 总 loss 时避免小数误差的标准写法，不能直接用sum(loss_list)累加 Tensor，会保留计算图引用导致内存泄漏。
第二个是device的统一。我见过太多新手在 CPU 上建好模型，又用model = model.cuda()迁移，但输入数据忘了.cuda()，然后报错说 device mismatch。最好的习惯是构造任何 Tensor 和网络前先定义统一的device，所有操作都通过.to(device)完成。把一个模型从单卡改成双卡或更多卡也是同样思路，不要用model.cuda()这种半迁移的方式，直接用model = nn.DataParallel(model)它会自己处理设备分配，但数据必须全部在默认设备上。
第三个是with torch.no_grad()的语义。验证阶段不需要梯度，如果不在这个上下文里跑 forward，autograd 会为验证过程构建额外的计算图，白白占掉一大块显存，而且理论上会把验证 loss 也带到反向传播图里。这个上下文管理器是绝对必须的。
3.4 损失函数和优化器的选择逻辑
CrossEntropyLoss在 PyTorch 里是“LogSoftmax + NLLLoss”的组合包。很多新手误以为它要求模型输出接一个 Softmax 层，其实恰恰相反：CrossEntropyLoss的输入应该是未经过 Softmax 的原始 logits，也就是网络的直接输出。如果在网络最后一层手动加Softmax，再用它做 loss，数值上虽然不会爆掉，但梯度计算会经历两次 Softmax 变换，训练速度会变慢，而且数值稳定性更好的是直接传 logits。这个坑我在很多项目代码里都见到过，属于标准的“看起来能跑、实际在走弯路”。
优化器选择上，Adam是绝大多数任务的首选起点，因为它自带自适应学习率，对初始学习率不那么敏感。SGD + Momentum则需要更好的学习率调度策略，通常在收敛精度上更高，但对调整经验要求也高。我的经验是：先用Adam配lr=1e-3跑通流程，如果后续需要刷精度，再切SGD（lr=0.1, momentum=0.9）并搭配余弦退火或 step decay。如果训练震荡厉害，就把 lr 调低一个数量级，不要动其他超参。
另外，lr_scheduler不要和手动调整optimizer.param_groups[0]['lr']混用，因为 scheduler 会按自己的节奏改写 lr，两者会互相覆盖。要么完全用 scheduler，要么完全不用，只在特定 epoch 手动设置，二选一。
3.5 断点续训：模型保存、加载与灾难恢复
模型训练动辄几十个小时，如果不做断点续训，机器重启或进程被杀就全部白跑。标准的保存策略是每轮 epoch 结束或 validation 指标提升时，保存一个 checkpoint，包含四样东西：模型状态字典model.state_dict()、优化器状态字典optimizer.state_dict()、当前 epoch、最佳指标值。
best_acc = 0.0 checkpoint = { "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "best_acc": best_acc, } torch.save(checkpoint, f"checkpoint_epoch{epoch}.pt") # 加载 checkpoint = torch.load("checkpoint_epoch10.pt") model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) start_epoch = checkpoint["epoch"] + 1这里两个常见误区需要记住。一个是不加载优化器状态，只看模型权重。这样做能恢复模型精度，但Adam的自适应学习率记忆和动量信息会丢失，后续训练前几个 step 会有一个明显的“冷启动期”，学习率调度也会错乱。另一个是torch.save和state_dict()的混用。torch.save(model, path)会序列化整个模型对象，包括模型结构定义，加载时容易出现类路径不匹配问题，尤其在项目重构后；state_dict()只保存参数数据，配合重新实例化模型再load_state_dict，这是更干净的方式。建议保存模型一律用state_dict()，仅在最终推理时可以用完整模型保存减少一步加载操作。
torch.load默认有个map_location参数，跨设备加载模型时很关键。在 GPU 上训练、到了只有 CPU 的机器上推理，加载时如果不指定map_location='cpu'，会报找不到 CUDA device 的错误。特指跨设备恢复 checkpoint 时，torch.load(path, map_location='cpu')是顺手一写但能省大麻烦的语法。
4. 训练过程的常见问题与排查技巧实录
4.1 显存爆炸（OOM）排查与优化
GPU 显存溢出是训练阶段发生率最高的问题之一。PyTorch 在显存不足时会抛出CUDA out of memory。我的常规排查步骤有三步：
第一步，看是不是真的显存不够，还是代码有泄漏。用nvidia-smi实时观察显存占用，正常一次迭代内存会增长然后回落，如果逐步爬升不回落，大概率是计算图没释放或 DataLoader 的num_workers线程没回收。另外，训练循环里如果每次迭代都累积了 Python 变量（比如把outputs或loss追加到一个列表里，而不是只取.item()），整个计算图的引用会一直保留，显存只增不减，这是新手最常踩的隐性坑。
第二步，降低单次送入显存的规模。最直接的是减小batch_size，顺带可以降输入图片分辨率或通道数。实际项目里，batch_size减半通常能让显存占用减半左右，但会牺牲一些梯度估计的稳定性。另一个常用技巧是给DataLoader开启pin_memory=True，虽然它本身不省显存，但能减少传输开销，间接缩短单次迭代时间，为调参争取空间。
第三步，用torch.cuda.empty_cache()和gc.collect()手动回收不需要的缓存块。要注意empty_cache()只是把 PyTorch 显存缓存里的空闲块返还给显存池，并不一定立即还给操作系统，但对后续分配确实有帮助。另外，开启梯度检查点torch.utils.checkpoint可以用计算换内存，在超大模型场景下才用得上，常规项目没必要。
4.2 损失变 NaN 的根因与处理
NaN 是另一个高频问题。最常见的原因有三个：学习率过大导致梯度爆炸、数据里有 NaN 或过大的值、损失函数输出为 0 后取 log 产生负无穷。
排查技巧是先看日志里 NaN 出现之前的最后一个 loss 值。如果 loss 在前几个 step 就直接 NaN，优先怀疑数据问题：检查输入是否有缺失值、标注是否越界（比如分类标签索引超出 num_classes）。如果 loss 是在某个 epoch 中途突然变 NaN，大概率是学习率过大导致梯度流中某个位置的梯度爆炸。快速验证方式是把 lr 降低 10 倍重跑，如果不再 NaN 就确认是 lr 问题。
处理梯度爆炸还有一个保底手段是梯度裁剪clip_grad_norm_，这是 Transformer 类模型训练的标准操作。设置max_norm=1.0表示把全局梯度的 L2 范数裁剪到 1.0 以内，防爆的同时基本不影响正常收敛路径。用torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)放在loss.backward()之后、optimizer.step()之前即可。
4.3 训练集表现好但验证集差的过拟合对策
如果训练 loss 持续下降、验证 loss 不降甚至上升，说明模型开始过拟合。常规解决方案按优先级排序：
- 增加数据增强（随机裁剪、翻转、旋转、颜色扰动），扩大有效样本空间；
- 增大 Dropout 概率或添加权重衰减（
weight_decay）； - 提前停止：监控验证指标，连续多个 epoch 不提升就恢复使用最佳 checkpoint；
- 降低模型容量：减少通道数或层数；
- 如果数据量很小，考虑用预训练模型做迁移学习而不是从零训练。
我个人特别推荐在每次验证指标提升时保存一次最佳权重，恢复后继续训练，这比“训练 100 轮然后考最后一个 epoch 的权重”稳得多。state_dict()保存和加载的代码在前面的断点续训部分已经给出了，直接复用即可。
5. 从单卡到多卡和面向生产的考虑
5.1 多 GPU 训练的两种路线
当单卡显存不够或训练时间太长，可以上多卡。PyTorch 提供了两条路线：一是nn.DataParallel（简称 DP），二是nn.DistributedDataParallel（简称 DDP）。
DP 的用法很简单，model = torch.nn.DataParallel(model)然后正常写循环就行。它的原理是把一个 batch 按 GPU 数量切分，每个 GPU 自己算一个子 batch 的梯度，再汇总到主卡更新参数。它的缺点也很明显：主卡承担了梯度汇总和参数广播的任务，容易成为瓶颈，而且主卡的显存占用明显高于其他卡，扩展性有限。
DDP 才是真正生产级的多卡方案。它的核心是多进程，每个进程持有一份模型副本，用 ring-allreduce 算法在训练过程中同步梯度，通信开销远小于 DP。DDP 的启动方式比较复杂，需要设置init_method、rank、world_size等环境变量。如果在单机多卡场景用torchrun --nproc_per_node=N train.py是最省事的启动方式，代码里只需要做这几件事：
import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # 初始化 dist.init_process_group(backend="nccl") rank = dist.get_rank() torch.cuda.set_device(rank) model = SimpleCNN().to(rank) model = DDP(model, device_ids=[rank]) # DataLoader 也要做分布式采样 from torch.utils.data.distributed import DistributedSampler sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, batch_size=64, sampler=sampler) # 每个 epoch 开始要设置 sampler 的 epoch，保证不同 epoch 顺序不同 sampler.set_epoch(epoch)DDP 的每一步都有讲究，backend="nccl"是 GPU 训练的标准后端，DistributedSampler保证每个进程分到不重叠的数据子集，set_epoch保证数据打乱顺序随 epoch 变化。
5.2 部署时的模型转换：torch.jit 与 ONNX
训练完成后要部署到服务端或边缘端，PyTorch 提供两条主流路径。一是用torch.jit.trace或torch.jit.script把模型编译为 TorchScript，这是一个可序列化、可优化、可脱离 Python 环境运行的格式。二是导出为 ONNX，再用各家推理引擎加载。
实际使用中，torch.jit.trace对大部分网络结构都能工作，它选一个代表性输入执行一次前向，把实际运行路径固化成静态图。但对于包含if分支或动态循环的模型，trace 只能固化一次执行路径，容易丢动态逻辑，这时需要用torch.jit.script，但 script 对 Python 语法支持有限，写法要保守。
导出 ONNX 的调用很简单：
dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "model.onnx", opset_version=12, input_names=["input"], output_names=["output"], )但 ONNX 导出有两个常见的坑。一个是某些操作（比如原地的inplace修改或不支持动态 shape 的计算）无法正确转换，建议导出前把model.eval()并关闭inplace操作。另一个是动态输入尺寸的问题，比如推理时输入长宽可变，需要在导出时显式标注动态维度，否则 ONNX 会固定输入 shape，部署后换尺寸就报错。
5.3 硬件加速与性能调优的粗浅经验
很多人忽略的是 PyTorch 本身的 benchmark 模式。在模型加载前执行torch.backends.cudnn.benchmark = True，CuDNN 会在第一次迭代时花更多时间做算法选型，之后使用最优算法，对固定输入 shape 的模型提升明显。如果输入 shape 变化频繁（比如目标检测中图片尺寸不定），这个开关可能反而带来不稳定，那就关掉。
另一个提升训练吞吐的手段是用torch.compile，PyTorch 2.x 引入的编译模式，通过图优化和算子融合，在 NVIDIA GPU 上对许多模型能带来 20%-40% 的训练速度提升。用法非常简单，把model = SimpleCNN()改成model = torch.compile(SimpleCNN())即可。但要注意torch.compile在实验阶段，某些自定义算子或动态控制流可能编译失败，生产环境上需要谨慎验证。
6. 形成属于自己的 PyTorch 工作流
说了这么多，最后总结一下我个人的建议路径。先不要急着刷各种模型的源码，先把张量操作、自动求导和训练循环这三块搞透，用一个小数据集（比如手写数字或 CIFAR-10）反复跑通。跑通之后，仔细去体会每个函数为什么存在，比如为什么zero_grad()必须在backward()之前，为什么验证要用model.eval()和no_grad()。
然后把三样东西练成本能：状态字典的正确保存与加载、统一的 device 管理、标准的.item()和.to()使用姿势。这些都是代码里最细碎的点，但决定你的代码能不能从“能跑”变成“能稳定跑”。
下一步再去接触主流模型库，读代码时你会发现自己已经能很流畅地理解一个模型的完整数据流。继续深入的方向包括：自定义损失函数和自定义层（继承nn.Module）、分布式的完整配置、ONNX 导出与推理引擎对接，以及基于 PyTorch 的机器学习运维经验，比如用mlflow或tensorboad记录实验指标。
我在实际项目里最深刻的体会是：PyTorch 的价值在于它的几个核心抽象（Tensor、autograd、Module、Dataset）彼此咬合得很紧，只要理解了这套抽象之间的关系，学任何基于它构建的上层工具都会快很多。反过来，如果只是抄代码，遇到 bug 时连排查的抓手都没有。所以建议多花点时间动手改改训练循环，改坏了大不了重来，但踩过的坑都会变成经验沉淀下来。如果最初跑模型时卡在哪一步，欢迎对照这篇文章里的排查顺序走一遍，大概率能解决你手头的问题。
