很长一段时间里，我身边的同事对深度学习框架的态度都是"能用就行"。直到有一次，在一个图像分割项目里需要排查一个非常诡异的梯度问题——模型在验证集上表现还算正常，可训练一两个 epoch 之后损失就突然变成 NaN。我用静态图框架排查了两天，日志打到手软,仍然找不到问题根源。后来换到 PyTorch，在出问题的那个算子后面加了一行 print，五分钟就定位到了原因。从那次之后，我就彻底倒向了 PyTorch，这几年从 CV 到 NLP 再到推荐系统，几乎所有项目都在用它。今天这篇内容，就是想以真实使用者的身份，聊聊 PyTorch 到底好在哪、坑在哪、怎么用才能少走弯路。
这篇文章适合两类人：一类是刚入门深度学习，想在 PyTorch 上建立起完整心智模型的新手；另一类是已经在用 PyTorch，但总觉得项目里差点意思、经常被各种报错折磨的进阶使用者。我会尽量避免教科书式的概念背诵，而是把项目里真正发生过的事情、真正解决问题和踩坑的过程写出来，希望能给你一些可以"抄作业"的参考。
1. 动态计算图：PyTorch 和"先编译后执行"派系最大的分水岭
1.1 计算图到底是什么？我用外卖订单来理解
很多教程一上来就讲计算图、自动求导、反向传播，概念一套接一套，初学者很容易被绕晕。我自己后来发现一个特别贴切的类比：计算图本质上就是一张"外卖订单流程图"。
想象你在餐厅点餐，订单会拆成"下单—备菜—炒制—打包—配送"几个步骤。深度学习模型的前向传播也是一样，输入数据流经一个个算子，每个算子（Convolution、ReLU、MatMul 等）就是流程中的一个环节。PyTorch 在每一次前向传播时，都会顺手把"这个环节的输入是什么、输出是什么、用了哪个算子"记录下来，形成一张动态的计算图。这张图保存了所有中间结果和相关操作的引用，目的只有一个：等反向传播时，能快速算出每个参数的梯度——就像外卖出问题了，能根据订单记录追查到是哪一步出了岔子。
PyTorch 的核心设计是动态计算图：每次前向传播都会重新构建这张图，图和代码是同步执行的，你代码里怎么写，计算图就是什么样。另一派采用"先编译后执行"的思路，需要先定义好整个模型结构，再整体编译成一张静态图，之后喂数据进来走这张固定的图。
这两者的差别在调试场景里体现得淋漓尽致。静态图模式下，你没法在模型中间随意插入一个 print 看中间张量的值，因为执行图是编译好的，数据和图是分离的。而 PyTorch 的动态图允许你在任何一行代码后面打印张量的 shape、dtype、数值分布，这在排查 NaN、shape 不匹配、梯度异常时简直是救命稻草。我说的那个图像分割项目，最后就是在某个归一化层后面打印出了 INF 值，一下定位到是输入里混入了异常像素，修复成本不到十分钟。
1.2 autograd 的细节：backward() 和 autograd.grad() 别乱用
PyTorch 的自动求导机制以autograd为核心，它记录的是张量之间的依赖关系。每个张量只要设置了requires_grad=True，它在参与运算后就会生成一个携带梯度函数的张量节点。调用loss.backward()时，PyTorch 会从 loss 出发，沿着计算图反向传播，把梯度填到每个叶子节点的.grad属性里。
这里有一个容易被忽略的点：默认情况下，一次 backward 之后计算图就会被释放。如果你在同一个计算图里调用两次backward()，第二次就会报错，提示你"计算图已经被释放"。这在某些场景下很让人头疼，比如你要计算两个不同 loss 对同一个参数的梯度，然后叠加。解决办法有两个：一是调用loss.backward(retain_graph=True)，保留计算图，但这样会显著增加内存开销；二是更推荐的做法，把两个 loss 直接相加，当成一个 loss 来 backward，这样只用构建一次计算图。
另一个实用细节是torch.autograd.grad()和backward()的区别。backward()把梯度累积到张量的.grad属性上，适合常规训练；grad()直接返回梯度值，不会修改.grad，适合需要临时获取梯度做分析的场景，比如可视化梯度分布、梯度裁剪时查看范数。这个区别在项目调试里非常有用，我在做梯度惩罚（Gradient Penalty）相关的对抗生成网络时，就经常用grad()单独取出某一层的梯度，避免影响正常训练流程。
还有一点，PyTorch 的with torch.no_grad():是一个高频使用的上下文管理器。它的作用是临时关闭自动求导，减少内存占用和计算开销。在模型评估、推理、计算评价指标时都应该加上它。很多人容易忽略的是，它也会影响 BatchNorm、Dropout 这些层的行为——不过那是模型训练/评估模式的问题，和no_grad本身无关，也要注意区分。
实操建议：如果你发现显存占用莫名其妙翻倍，十有八九是某个地方没加with torch.no_grad()，或者没有用detach()切断某个张量的梯度传播。
2. 张量理解与 nn.Module：为什么 PyTorch 写模型像搭积木
2.1 张量的心智模型：先想清楚 shape、device、dtype 三件事
很多初学者在写 PyTorch 代码时，遇到shape mismatch报错就懵了。我想说，张量报错里 90% 都出在三件事没想清楚：shape（形状）、device（设备）、dtype（数据类型）。
我给你一个排查 shape 问题的固定思路：写任何一层网络之前，先在草稿纸上写下输入张量的 shape，然后逐个算子手算输出 shape。比如输入是(B, C, H, W)，过一个 3x3 卷积、padding 为 1、stride 为 1 的层，输出 shape 依然是(B, C', H, W)；过一个 stride=2 的池化，高宽各自减半。卷积层和全连接层的接口参数不同，全连接层要求输入是二维(B, D)，所以卷积输出要flatten成二维再喂进去，这是最常见的报错原因之一。
device 问题也经常让人崩溃。最典型的是：模型参数在 GPU 上，输入数据在 CPU 上，一跑就报错“Expected all tensors to be on the same device”。我习惯的做法是写一个device = torch.device("cuda" if torch.cuda.is_available() else "cpu")，然后模型和数据都显式.to(device)，绝不依赖某个默认设备。多卡环境下还需要注意torch.cuda.set_device，否则可能会出现张量在cuda:0、模型在cuda:1的错位。
dtype 的问题常见于混合精度或某些算子对数据类型的限制。比如某些算子只支持 float32 而不支持 float64；torch.matmul在半精度下可能出现精度下降。此外，输入图像如果不小心转成了 uint8，很多算子在autograd下就没法求导了。统一用torch.float32基本是最稳妥的选择，混合精度时再用autocast自动管理。
2.2 nn.Module 内部机制：参数、缓冲区、子模块的注册
nn.Module是 PyTorch 所有模型层的基类，它的设计哲学用一个词概括就是"递归式参数管理"。你定义了一个模型，它内部可能有若干子层，每个子层又有自己的参数，nn.Module通过特定顺序的注册机制把这些参数收集起来。理解了这个机制，你就理解了为什么model.parameters()、model.state_dict()能自动拿到所有层里的可学习参数。
这里面最关键的是parameter 和 buffer 的区别。可学习参数必须注册为nn.Parameter，比如self.weight = nn.Parameter(torch.randn(...))，这样它才会被优化器识别。而 BatchNorm 里的 running_mean、running_var 这类不属于梯度更新范围但需要被记录和张持久化的量，应注册为 buffer：self.register_buffer("running_mean", torch.zeros(...))。如果你把 Buffer 误注册成 Parameter，训练时会试图更新它，但求导又求不出来，随后报错或数值混乱都可能出现。
forward()方法的写法也有讲究：你定义了self.conv = nn.Conv2d(...)，然后调用model(x)时，真正执行的是__call__方法，它会先检查钩子、再调用你写的forward(x)返回结果。这也是为什么我们不能直接调用model.forward(x)——那样会跳过__call__里的各种机制，比如nn.Module的 hooks、训练/评估模式的切换逻辑都会失效。
自定义一个层时，顺带提醒一个心得：尽量在__init__里把所有参数都定义好，不要在forward里新建nn.Parameter。在forward里动态创建参数，虽然某些极端场景下能跑，但会破坏 PyTorch 对参数注册的直觉，分布式训练和模型保存时容易出现"参数突然消失"的诡异问题。
2.3 用积木思维组织模型
我建议把所有模型都按"积木"来组织：底层是nn.Conv2d、nn.Linear、nn.LayerNorm之类的原子积木；中间层用nn.Sequential或自定义Module组装成残差块、Transformer 编码器、特征提取器等语义模块；最顶层用一个大Module组合这些模块，提供forward逻辑。
我见过有些人习惯在一整个巨型nn.Module里写上百行forward，所有层全部塞在__init__里。项目初期没问题，等后面要改结构、做消融实验或者排查速度瓶颈时，就会非常痛苦。更好的做法是像写工程代码一样拆分模块，每个模块只做一件事，用nn.Sequential处理线性的栈式结构，用自定义Module处理有分支或跳跃连接的复杂结构。这样不管是单机调试还是后续迁移到分布式训练框架，都能省一大半力气。
3. 数据管线与训练循环：被低估的工程重心
3.1 DataLoader 的 num_workers 和 pin_memory 决定训练速度的天花板
模型结构写得再快，数据跟不上也是白搭。很多人抱怨训练慢，一上来就怪 GPU 不够好，其实很多时候瓶颈在数据加载这一环。torch.utils.data.DataLoader里有几个参数，每个都值得你仔细调教。
第一个是num_workers。它决定了用几个子进程来做数据预处理和加载。默认值是 0，意思是数据加载在主进程里同步进行，GPU 每次要算完一个 batch 就得干等数据。设置为 2、4、8 甚至更高，可以让数据加载和 GPU 计算并行起来。但这个值不是越大越好：开太多 worker 会带来进程间通信和内存争抢的开销，反而拖慢速度。一般可以先用num_workers = 4起步，逐步观察 GPU 利用率来优化。
第二个是pin_memory=True。这个参数会锁页内存，让数据从 CPU 到 GPU 的拷贝走更快的内存通道，实测在很多数据集上能带来 10%-20% 的提速。代价是占用更多 CPU 内存，如果内存紧张就要权衡。
第三个是persistent_workers=True。DataLoader 默认每个 epoch 结束后会销毁 worker 进程，下个 epoch 再重新创建，这个创建开销其实不小。设置persistent_workers=True可以让 worker 长期驻留，规避反复创建的损耗。实测在数据集较小、epoch 数量很多的情况下，收益非常明显。
我还想特别提醒一下：如果你的数据增强逻辑里包含大量 CPU 计算（比如随机裁剪、颜色抖动、缩放等），而num_workers又不够，那么瓶颈一定会出现在数据增强上。解决思路是尽量在__getitem__里做必要的最小集处理，复杂的增广可以放到 GPU 上用张量操作完成，或者直接用现成的数据增强库优化。
3.2 一个标准的训练循环骨架
训练循环代码虽然看着简单，但工程上要注意的细节远比想象的多。我给出一个我习惯使用的骨架：
for epoch in range(num_epochs): model.train() train_loss = 0.0 for batch_idx, (inputs, labels) in enumerate(train_loader): inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() # 梯度裁剪，稳定训练 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() if batch_idx % log_interval == 0: print(f"Epoch {epoch} Batch {batch_idx} Loss {loss.item():.4f}") # 每个 epoch 结束评估 model.eval() val_loss = 0.0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) val_loss += criterion(outputs, labels).item() print(f"Epoch {epoch} Val Loss {val_loss / len(val_loader):.4f}") # 保存 checkpoint这个骨架里有几个容易踩的细节。optimizer.zero_grad()必须在每个 batch 开始前清空梯度；否则梯度会跨 batch 累积，导致训练不稳定。model.train()和model.eval()的切换必须及时：train()模式下 BatchNorm 会利用当前 batch 统计量更新 running 统计量，Dropout 会随机丢弃神经元；eval()模式下两者都会关闭，只做确定性推理。很多新手在评估时忘记model.eval()，验证集指标会飘忽不定，这其实是 Dropout 或 BatchNorm 行为没有被正确锁定导致的。
还有 checkpoint 保存。我强烈建议按"epoch 数 + 验证指标"命名，比如ckpt_epoch10_val_loss0.23.pth，并同时保存model.state_dict()、optimizer.state_dict()、epoch以及一些超参数。这样不仅能断点续训，还能方便回溯最优模型。曾经我见过只保存model.state_dict()的项目，回滚到某个 epoch 时优化器状态丢了，学习率调度也乱了，训练半天效果始终不如保存过完整状态的版本。
3.3 学习率调度、梯度累积与 EMA：三个稳定训练的实用技巧
学习率调度是我非常关注的一环。PyTorch 内置了torch.optim.lr_scheduler全家桶，包括StepLR、CosineAnnealingLR、ReduceLROnPlateau等。我的经验是：先用ReduceLROnPlateau跑一遍，让指标变化告诉你什么时候该降学习率，而不是凭空拍一个下降节点。很多项目在跑 NLP 预训练或大规模视觉模型时，会搭配get_cosine_schedule_with_warmup这类 warmup + 余弦退火的组合，效果更稳定。
梯度累积是一种显存不够时的常用手段。把optimizer.step()改成每累积 N 个 batch 才执行一次，本质上就是用时间换空间，让你能模拟更大的 batch size。写法上要额外注意：累积阶段要用loss.backward()累积梯度，但不要step()和zero_grad()；达到 N 个 batch 后再 step 并清零。
EMA（指数移动平均）则是对模型参数做平滑：维护一份参数的滑动平均副本，评估时用这份副本而不是实时参数。动手实现其实很简单，就是每个 step 更新一下：
ema_param = decay * ema_param + (1 - decay) * model_param经验值是 decay 取 0.999，每 step 更新一次。使用 EMA 后，很多模型的泛化能力都有肉眼可见的提升，尤其是训练后期 loss 震荡明显时，EMA 版本的效果曲线会平滑很多。
4. 我在 PyTorch 项目里反复踩过的四个坑
4.1 inplace 操作与 autograd 的爱恨情仇
inplace操作指的是像tensor += 1、tensor.relu_()这种直接修改原张量值的操作。它们在节省内存方面很有用，但和自动求导机制存在天然矛盾。PyTorch 的反向传播依赖前向传播时保存的中间结果，如果你对某个需要求导的张量做了 inplace 修改，原始值就丢了，反向传播时梯度算不准，甚至直接崩溃。
最典型的案例是ReLU的 inplace 版本。很多基准代码里会出现nn.ReLU(inplace=True)，这是为了省显存。但在一些复杂的网络分支结构里，同一个张量被多个分支引用，其中一个分支做了 inplace 修改，另一个分支在反向传播时拿到的已经是修改后的值，梯度就会出错。我的经验是：小模型完全不用开 inplace；大模型实在要开，必须保证该张量没有被其他分支引用。
排查 inplace 报错还有一个技巧：报错信息通常会提示one of the variables needed for gradient computation has been modified by an inplace operation。顺着堆栈找到出问题的张量，看看它被哪些操作引用了，最有效的办法是把相关属性的 inplace 参数全部关掉，再跑一次，通常就不会报错了。虽然显存稍微多占一点，但换来的是安心。
4.2 模型和数据被反复在 CPU 和 GPU 之间搬运
这个问题非常隐蔽。有时你会写类似loss.item()这样的代码，它会把 GPU 上的标量搬到 CPU 上；如果某个循环里频繁调用这个操作，就会产生大量的同步开销，让 GPU 的利用率一降再降。更隐蔽的是，在DataLoader的collate_fn里不小心用了 CPU 上的 numpy 操作，数据被搬回 CPU，后面又被迫搬回 GPU，一来一回就是一次不小的开销。
我的建议是：在训练主循环里，尽量减少.item()、.cpu()、.numpy()的频率。如果确实需要记录 loss 曲线，可以每个 log_interval 才取一次，不要每个 batch 都取。另外要注意，某个 batch 里的张量如果被强制.cpu()了，后续再调用.to(device)又是一个拷贝，这种无意义的搬运在生产代码里特别常见。
4.3 随机种子、可复现性与分布式训练的坑
深度学习项目需要随机性，但复现实验时必须保证随机种子一致。PyTorch 的随机性来源有三处：PyTorch 本身的随机算子、Python 的random模块、numpy 的随机数生成器。你只设置torch.manual_seed(42)是不够的，因为数据增强里大量使用 numpy 的随机数，如果 numpy 的种子不固定，实验结果依然不可复现。
在数据加载阶段，DataLoader的 worker 进程会各自继承主进程的随机状态，如果你不额外给每个 worker 设置独立的种子，多 epoch 训练时可能会出现每个 epoch 数据顺序完全相同的情况——这在某些场景下不是好事。PyTorch 提供了torch.utils.data.DataLoader的generator参数，配合自定义 worker_init_fn 可以很好地控制。多卡分布式训练时还有一点：如果不用DistributedSampler，每个进程会看到同一份数据，相当于把同一个 batch 重复训练了 N 次，梯度更新方向会发生偏差，模型效果会莫名其妙变差。
4.4 state_dict 键名不匹配：加载预训练权重最常见的报错
我们在做迁移学习或加载别人开源的预训练模型时，最常遇到的就是UnExpected key(s) in state_dict或者Missing key(s) in state_dict。前者说明模型里没有某些键，后者说明加载时没匹配上所有键。
遇到这种情况不要慌，先打印一下两边的键名对照看看。通常原因有几种：一是新旧版本的模型结构不同，比如把conv1改成了stem_conv；二是state_dict的键名带上了module.前缀，这是多卡训练保存下来的常见现象；三是分类头名称不一样，导致最后一个全连接层对不上。
我之前在某图像分类项目中加载预训练权重时，发现只有分类头那几层的键不匹配。处理办法很简单，加载时加上strict=False，这样除了分类头之外的层都会加载成功，分类头保留随机初始化。千万不要把整个 state_dict 的键名硬改成和模型完全一致，那样反而容易覆盖掉正确的预训练信息。
5. 性能优化：从单卡到多卡，让显存和吞吐都体面
5.1 显存占用分析：别再用 nvidia-smi 瞎猜了
很多人的显存优化方法就是盯着nvidia-smi反复看，其实这个工具只能告诉你总显存占用，没法告诉你到底是哪些张量占了空间。PyTorch 提供了两个非常好用的工具：torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。前者是真正分配出去且被张量占用的内存，后者是 PyTorch 向 CUDA 申请但尚未用完的内存。
想逐层分析显存占用，可以写一个简单的 hook，或者使用 PyTorch 的torch.cuda.memory._snapshot()API。但我的经验是，真正做显存优化时，先做一个简单的判断：显存瓶颈是在前向传播还是反向传播？前向会保存中间激活值用于反向传播计算梯度；反向会额外计算每个参数的梯度。如果你发现中间激活值占用特别大，可以优先考虑使用更小的 batch size、开启梯度检查点（gradient checkpointing）、或者用混合精度。
5.2 混合精度训练：用半个精度换一倍显存和接近一倍的吞吐
现代 GPU（特别是带 Tensor Core 的）对半精度计算支持非常好，混合精度训练目前已经是大型项目的基础配置了。PyTorch 从 1.6 开始提供了torch.cuda.amp，用法非常简单：
scaler = torch.cuda.amp.GradScaler() for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()关键点在于GradScaler。半精度能表达的数值范围很小，梯度在反向传播时容易下溢到 0，所以需要先把 loss 放大若干倍，完成反向传播后再缩回来。GradScaler做的就是这件事。很多人第一次用 AMP 时报"loss is NaN"或"gradient is None"，往往是因为忘记用GradScaler，或者loss.backward()用在了autocast上下文之外。
使用混合精度后，模型权重仍然以 float32 保存，只有计算过程和中间激活值使用 float16。所以模型精度通常不会明显下降，但训练速度和显存占用都会有质的改善。实测在较大规模视觉模型上，AMP 可以把显存峰值降低 40% 左右，吞吐提升 40%-80%，几乎是性价比最高的优化手段。
5.3 梯度检查点与梯度累积：空间换时间与时间换空间
显存不够的第二个大杀器是梯度检查点（Gradient Checkpointing）。它的原理很聪明：训练时原本要保存所有中间激活值，梯度检查点则只保存其中的一部分关键节点；反向传播时需要用到没保存的中间值，就重新做一次前向传播来"重算"它们。这样显存占用显著下降，代价是额外的计算量。用起来也很简单，PyTorch 提供了torch.utils.checkpoint.checkpoint_sequential和torch.utils.checkpoint.checkpoint接口。通常我们把 checkpoint 用在几个占用显存特别大的大模块上，而不是逐层应用，这样只需多花大概 30% 的计算量就能省下 60% 以上的显存。
梯度累积我在前面已经提过，实际上是把多个 batch 的梯度攒起来一次性更新。它和增大 batch size 并不完全等价，因为 BatchNorm 的统计量仍然基于单个 batch 计算，但训练稳定性已经能应付绝大多数场景。显存极度紧张时，4 卡梯度累积 × 真实 batch size 缩小 4 倍的效果，和单卡大 batch 的效果差别并不大。
5.4 多卡训练选型：数据并行、分布式数据并行、还是模型并行？
如果你有多张 GPU，最简单的是torch.nn.DataParallel，它把每个 batch 切成多份分给各卡计算梯度，然后汇总到主卡更新参数。用法极其简单，model = nn.DataParallel(model)一行搞定。但它的缺点是主卡负载很高，多卡扩展性不佳，不建议大规模使用。
更推荐torch.nn.parallel.DistributedDataParallel（简称 DDP）。它每个进程持有完整模型副本，各处理各的 batch，通过梯度通信同步参数。虽然配置起来要写一些初始化代码（init_process_group、DistributedSampler等），但性能和扩展性都远胜 DataParallel。我的经验是：只要超过一张卡，就直接上 DDP，不要再用 DataParallel。第一次配 DDP 后会有一个适应期，但是一旦配好，后面换机器、加卡都非常顺利。
如果你的模型大到单卡放不下，就需要模型并行了。PyTorch 提供了torch.distributed.pipeline.sync.Pipe这类工具，可以把模型的不同层放在不同的 GPU 上。但这是相对高级的玩法，普通项目一般用不到；真到了这一步，你大概率已经对显存和通信非常熟悉了，可以自己设计流水线。
6. 训练完了不是终点：PyTorch 模型的导出与部署
6.1 torch.jit.trace 和 torch.jit.script，先理解再选择
训练完模型后，直接保存model.state_dict()只是训练侧的做法。部署场景下，我们要的是"模型可独立运行的结构+权重"整体文件，而且最好脱离 PyTorch 运行库也能被推理引擎解析。PyTorch 官方提供的 TorchScript 是第一步。
torch.jit.trace是"跟踪"模式：喂一个示例输入进去，记录所有算子执行路径，生成一个静态脚本化模型。它实现简单，但缺点明显——遇到条件分支、循环等动态逻辑会失效，因为跟踪过程只走了一条路径。torch.jit.script则是"脚本化"模式，会真实解析你的 Python 代码，支持动态控制流。代价是它对代码语法有一定要求，需要兼容 TorchScript 的语法子集。
我的选型经验是：能用 trace 就用 trace，不能用就用 script。很多模型的 forward 里只要不包含运行时依赖输入的 if/else，trace 完全够用，生成速度快、兼容性好。但如果你用过自定义的 Python 数据结构或者复杂的控制流，那就要上 script，或者干脆改写代码，让 trace 变成可行。
6.2 ONNX 导出：跨平台推理的中间桥梁
在很多生产环境里，模型不一定要用 PyTorch 直接跑，而是转成 ONNX 格式交给推理引擎处理。PyTorch 提供了torch.onnx.export接口，核心就是指定模型、示例输入、输入输出的名称和动态轴。
torch.onnx.export( model, dummy_input, "model.onnx", opset_version=13, input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}, )这里dynamic_axes很关键：如果不声明 batch 维度是动态的，导出后的 ONNX 模型会把 batch size 固定死，线上服务一旦换了一个 batch size 就报错。导出这一步真正的麻烦是算子的兼容性。PyTorch 里有几百种算子，ONNX 的算子集不是每个都支持。特别是某些自定义层、torch.where、torch.topk在某些 opset 版本下会出现不兼容；遇到这种情况，常见的解决办法是把模型拆成几个部分分别导出，或者把不兼容的算子替换为等价的基础算子组合。
导出之后，记得用onnxruntime或者推理引擎自带的工具跑一遍，核对输出和 PyTorch 原始输出的数值误差。我遇到过很多次这种情况：PyTorch 单测上完全一致，但推理引擎里由于算子实现、精度模式的差异，输出有细微偏差。所以要设置一个合理的容差阈值，否则下游任务会受到不可预期的干扰。
6.3 大规模服务化：把模型变成可对外提供能力的服务
到了真正提供服务时，常见的部署方案有几种：用 TorchServe、用torch.jit打包后用纯 C++ API 加载运行，或者转 ONNX 后配合推理引擎。到底选哪种，取决于你的推理延迟要求和工程团队的技术栈。
如果是快速原型或对延迟要求不高，直接用 FastAPI 套 PyTorch 推理就能跑起来：加载模型、预处理输入、前向、后处理返回结果。这个方案实现成本最低，但每个请求都走 Python 进程，并发一高就容易成为瓶颈。如果对吞吐和延迟要求苛刻，我一般建议转 ONNX 并由专门的推理引擎托管。很多团队还会配合模型量化（如动态量化、静态量化），进一步压缩模型体积和延迟。
这里要提一个非常容易被忽略的工程坑：推理时的预处理必须和训练时完全一致。图像缩放算法、归一化均值方差、通道顺序，任何一个不一致都会导致模型效果崩掉。我在某个项目里排查过线上推理异常，最后发现是线上服务的图像缩放算法和训练时差了一个像素的插值方式。这类问题不报错，但直接影响业务指标，排查起来很费时。
6.4 最后聊一点部署心得
如果非要说一个大原则，我想是：训练和部署的代码要尽量共用一套数据预处理逻辑。把预处理函数单独抽出来，训练时调用它，部署时也调用它，从根源上杜绝不一致。
还有一个小技巧，导出的模型一定要做"回归测试"：拿一批训练集里没见过的样本，分别跑 PyTorch 原模型和导出模型，对比数值差异是否在容忍范围内。这个测试在每次更新模型后都要跑一遍，不要偷懒。我见过导出时因为某个算子不兼容被偷偷替换成近似算子，单看一个样本输出没问题，一压测就发现某些输入下误差放大，最后花了一整天追查，才发现是导出的坑。
坏消息是，部署环节的坑远不止这些；好消息是，只要把预处理、导出回归测试、精度检查这几个要点做好，大多数问题都能在早期暴露，不至于到线上才翻车。
在我自己的项目里，PyTorch 几乎成了默认选项。它学习曲线相对平缓、调试方便、社区资料丰富，能覆盖从科研原型到生产部署的整个链路。虽然它也有自己的缺点——比如多进程训练的内存占用偏高、TorchScript 对新特性的支持不够及时——但相比它带来的研发效率，这些小问题完全可以通过合理的工程习惯来规避。希望这些实际踩坑得来的经验，能帮你少走我走过的那些弯路。
