PyTorch 模型搭建核心：从环境配置到训练导出实战

我最早接触 PyTorch 时的第一印象不是官网首页那句漂亮的 slogan而是一条绕不开的报错RuntimeError: CUDA out of memory。真实情况是我连网络还没搭好双卡机器上一张卡被别的任务占满另一张卡的显存也给模型预留少了。后来把环境、显存、batch_size、模型大小全部算了一遍才明白很多时候问题根本不是“模型的 forward 写错了”而是“这层模型在 PyTorch 里的运行链路”没理顺。这篇文章不打算做那种一步步点击官方文档的保姆级翻译而是想站在“我要真正上手搭一个模型、跑通一轮训练、再把它导出部署”的角度把 PyTorch 模型搭建核心与基本使用方法梳理一遍。适合两类人一类是刚装完 torch正被各种版本、设备、张量报错打得烦躁的初学者另一类是已经能跑通简单案例但想更系统地读懂nn.Module、训练循环、序列建模和导出流程的开发者。你不需要记一堆冷门 API只需要把这里面的几条核心链路打通后面遇到什么模型都能用同一套思维快速接住。1. 环境与版本对应为什么你装完 torch 第一步就跑不通很多人学 PyTorch 的第一个挫败感不是来自模型写不出来而是来自“我明明照着官网复制了安装命令怎么一运行就报错”。这类问题十有八九出在环境隔离、CPU/GPU 版本混选、Python 和 torch 版本不匹配这三件事上。这三个问题如果在开头没处理好后面每写一段代码都可能被环境问题反复打断。1.1 先建一个干净的 conda 环境别让全局 Python 替你背锅我见过有人在系统自带 Python 里直接pip install torch之后装其他依赖时动不动就冲突最后整个 Python 环境被搞坏。PyTorch 对底层依赖非常挑剔尤其是numpy、protobuf、typing_extensions这些库版本差一两个小版本都可能触发匪夷所思的错误。此时 conda 或 venv 的隔离环境不是可选优化而是必需手段。我最常用的是 condaconda create -n pytorch python3.10 conda activate pytorch创建好后所有后续安装都进这个环境。这里补一个教训如果你不确定用 conda 还是 pip 装 torch我的建议是——用 pip 就从官网提供的 pip 命令装用 conda 就全程 conda 装尽量不要混着来。conda install pytorch -c pytorch和pip install torch两套体系虽然最终都装进 site-packages但过程中可能触发不同版本的numpy、mkl、cudatoolkit重装来回摇摆很容易收到一堆依赖冲突警告。1.2 CPU 版与 GPU 版到底差在哪怎么选不少人一上来就想装 GPU 版但实际上有没有 NVIDIA 显卡决定了你能装什么。没有独立显卡或只有核显安装 CPU 版完全没问题模型能跑只是慢有 NVIDIA 显卡才需要考虑 GPU 版。GPU 版本质上是把 PyTorch 的核心算子编译成调用 CUDA 运行库的版本使得张量可以放到显存中计算。官网首页给的是这种命令行pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121这里cu121表示 CUDA 12.1 运行时版本。需要注意绝大多数用户并不需要手动装完整的 CUDA Toolkit因为 PyTorch 的 wheel 包已经打包了它依赖的 CUDA 运行库。你需要确保的只有一件事显卡驱动版本够新足以支持这个 CUDA 运行库。查驱动是否匹配可以用nvidia-smi它右上角会有一行CUDA Version表示当前驱动最高支持的 CUDA 版本只要这个数字大于等于 PyTorch 要求的运行库版本通常就没问题。场景推荐安装方式验证命令只有 CPUpip install torch --index-url .../cputorch.cuda.is_available()应为 FalseNVIDIA GPU按驱动支持的 CUDA 选cu118/cu121等torch.cuda.is_available()应为 TrueApple Silicon官方提供的 macOS 版本torch.backends.mps.is_available()装完后先别急着跑模型花十秒钟做一次体检import torch print(torch.__version__) print(torch.version.cuda) print(torch.cuda.is_available()) if torch.cuda.is_available(): print(torch.cuda.get_device_name(0))如果torch.cuda.is_available()返回 False不要急着重装先把驱动版本、wheel 的 cu 编号、Python 位数这三项逐一核对。我踩过的经验是很多时候不是 torch 装错了而是 Windows 下明明有显卡驱动但 Python 进程运行在某个禁用 GPU 的远程会话里或者驱动是几天前刚被 Windows Update 覆盖成旧版。1.3 Python 版本与 torch 版本怎么对应PyTorch 每个版本都有对应的 Python 版本范围官方 release note 里写得很清楚但新手一般不会专门去查。如果硬要给一个不会出错的建议新建 conda 环境时选 Python 3.10 或 3.11装 torch 2.x 的近期稳定版本。这套组合是社区里覆盖面最广、第三库兼容性最好的方案。老项目用的 Python 3.8、3.9也基本能装 1.x 和部分 2.x 的 torch但没必要刻意去追最新 Python。Python 3.12 或 3.13 用户要看 PyTorch 官方是否已经提供对应 wheel很多自定义算子编译库可能还停留在旧版本。还有一个很容易被忽略的问题torch、torchvision、torchaudio这三个包必须配套安装版本号要来自同一次发布的组合。比如你把 torch 升级到 2.2.2却只装了 torchvision 0.17.0很多接口可能仍然能跑但某些图像变换和新模型定义会出现“torchvision 版本过低”的隐性问题。官网get-started页面给出的命令三者版本就是捆绑一致的别拆开乱装。1.4 torchvision、torchaudio 是生态配套别当成无关紧要的东西有些教程为了省时间只装 torch。等你想加载torchvision.models.resnet18或者用torchvision.transforms做图像预处理时才发现缺包。我的建议是只要是你可能用到图像、音频、视频数据就直接把三个包一起装。它们与 torch 版本强绑定前面表格里的推荐命令都是一个整体。安装完成后用python -c import torchvision; print(torchvision.__version__)检查一次确保能正常导入。2. 张量与自动求导PyTorch 一切的底层逻辑在这里环境准备好了接下来要面对的是 PyTorch 最核心的抽象张量和自动求导。这两个概念理解不到位后面写模型时你会到处靠猜猜不中就得反复调试。理解到位了会发现所有模型结构本质上都是“张量之间的运算图”。2.1 张量是 NumPy 数组的“深度版”但别用 NumPy 习惯写它PyTorch 里的torch.Tensor在直觉上非常接近 NumPy 的ndarray有shape、dtype、索引、切片、广播规则。但有两处不同非常关键张量可以放在 GPU 显存上通过.to(cuda)在 CPU/GPU 之间迁移张量可以挂在计算图上通过requires_gradTrue追踪梯度。先记住最基本的创建方式import torch a torch.tensor([1.0, 2.0, 3.0]) # 从列表创建 b torch.zeros(2, 3) # 全零 c torch.randn(2, 3) # 标准正态随机 d torch.arange(0, 10, 2) # 等差数列 e a.reshape(3, 1) # 改变形状这里有个特别容易混淆的细节torch.tensor()是工厂函数torch.Tensor()也可以用来创建张量但二者并不等价。torch.Tensor()不指定参数时创建的是空张量而且它默认 dtype 是torch.float32而torch.tensor([1, 2, 3])会保留整数类型。新手如果拿这两种方式交替用很容易在处理图像标签时踩到 dtype 不一致的坑。2.2 requires_grad 背后的 reverse-mode autodiff当我第一次看到loss.backward()时最大的疑惑是PyTorch 怎么知道对谁求导答案是你创建的每个张量以及基于它做的每次运算都会被记录进一张动态计算图。张量有个属性叫grad_fn记录它是通过什么运算生成的没有grad_fn的原始输入张量称为叶子张量。一个最简单的验证x torch.tensor(2.0, requires_gradTrue) y x ** 3 z y.sum() z.backward() print(x.grad) # 12.0因为 dz/dx 3*x^2 12反向传播执行时PyTorch 从z出发沿grad_fn链路反向走回叶子节点x把梯度写入x.grad。这就是 reverse-mode autodiff。理解这一点后你就能明白为什么optimizer.zero_grad()那么重要如果不清空上一次梯度x.grad会不断累加。梯度累加在某些特殊训练策略里是有用的东西但它绝不是大多数情况的默认预期。关于参数优化器那一层你只需要记住定义网络时把需要更新的权重用nn.Parameter包起来优化器拿到model.parameters()每走完一个 batch 就基于.grad更新参数。模型搭建的核心其实就是管理这些Parameter的集合。2.3 no_grad 与 detach推理和特征提取的正确姿势训练时需要梯度推理或验证时不需要。此时有两种常见做法包在with torch.no_grad():下或者用.detach()从计算图中脱离。二者区别简单说torch.no_grad()是一个上下文管理器这里的所有运算都不会被记录到计算图省显存、省时间.detach()是返回一个新的张量它与原张量共享底层数据但requires_grad为 False。我在验证集上计算准确率时永远会写model.eval() with torch.no_grad(): outputs model(data) preds outputs.argmax(dim1)如果不加no_grad仅仅验证一次准确率也可能累积一整张计算图显存占用一路攀升。这听起来是小事但很多人遇到“训练时显存慢慢涨到爆”往往就是在这里漏了。2.4 设备一致性CPU/GPU 张量的红线“Expected all tensors to be on the same device” 这类报错是 PyTorch 新手会遇到的几乎最为普遍的问题。原因是CPU 上的张量与 GPU 上的张量不能直接做四则运算。解决办法也很固定一致地调用.to(device)。经验是在一开始就定义device torch.device(cuda if torch.cuda.is_available() else cpu)然后把模型、输入数据、标签全部统一迁移到device上。不要一会儿.cpu()一会儿.cuda()那样代码写起来费劲还容易漏掉某些变量。模型里的参数注册好后直接model.to(device)一次性迁移即可。3. nn.Module 是模型搭建的骨架理解这几个内部机制才能自由扩展nn.Module是 PyTorch 模型搭建核心的核心。很多人以为它只是一个“容器”把层塞进去、定义forward就能用。实际上它做了三件非常重要、但默认不可见的事参数注册、设备迁移、训练/评估状态切换。这些机制不理解你可能很难解释为什么有些写法能训练有些写法会把某个层忘在 CPU 上。3.1__init__注册机制不是约定而是行为import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 256) self.fc2 nn.Linear(256, 10) def forward(self, x): x self.fc1(x) x nn.ReLU()(x) # 不推荐见下文 return self.fc2(x)上面这段代码能跑但你注意到没有nn.ReLU()写在forward里每次都会创建新实例而且它没有参数所以影响不大。真正的问题是一些带参数的层如果被写在forward里而不是__init__里它的参数不会被model.parameters()捕获训练时自然也不会更新。这是因为 PyTorch 是在__init__过程中扫描self属性并递归注册的。想判断一个层有没有被注册最简单的办法是打印modelMyNet( (fc1): Linear(in_features784, out_features256, biasTrue) (fc2): Linear(in_features256, out_features10, biasTrue) )如果某个层没有出现在这个结构里那它在state_dict和.to(device)面前基本是透明状态后面会莫名其妙。3.2 三种构建网络的姿势继承类、Sequential 与 ModuleList/Dict不是所有网络都必须继承nn.Module并手写forward。根据网络结构复杂度一般有三种姿势第一种手写nn.Module子类。这是最灵活的方式适合需要定制 forward 逻辑的模型也是绝大多数研究模型的写法。第二种nn.Sequential。适合那些结构固定、一层接一层、不带多分支的网络比如一个小 MLPmodel nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10), )Sequential本身也是一个Module它的forward就是按顺序执行子模块。简单网络用它最省事。第三种nn.ModuleList和nn.ModuleDict。这两种容器解决的是“层数量动态变化”或“按 key 组织模块”的问题。比如你要写一个 7 层还是 12 层由配置文件决定的 ResNet 风格网络用 Python 列表直接存nn.Module不会被注册换成nn.ModuleList就行。layers nn.ModuleList([ nn.Linear(64, 64), nn.Linear(64, 64), nn.Linear(64, 10), ]) for layer in layers: x layer(x)如果以后看到代码里用普通list存模块训练时参数不更新先怀疑这里。3.3 为什么不直接调 forward而是调 model(x)model(x)和model.forward(x)在简单情况下结果一样但行为有本质区别。__call__内部会先执行一些 hooks再走到forward。这些 hooks 是 PyTorch 高级功能的入口比如打印中间层输出、剪枝、权重可视化。只要你在网上看到“注册 forward hook”它依赖的就是这条__call__链路。所以除非你清楚自己在做什么否则永远通过model(x)调用网络。很多人为了方便调试直接调model.forward(x)结果特征提取钩子全部失效排查到天亮。另外大量官方代码里面的麻烦报错都指向“在forward里用了原地修改操作”比如x 1。因为自动求导需要记录原始值原地修改会破坏计算图。你要是发现报错信息里出现inplace operation字样去检查 forward 里的、*、nn.ReLU(inplaceTrue)这些操作即可。3.4 打印模型与返回实例的类对象名称排查结构的实用技巧另一个热度很高的问题是如何在代码里获取模型实例的类对象名称。打印整个model时显示内容是由__repr__控制的它包含类名和子模块结构。如果只是想拿到类名字符串用type(model).__name__或者model.__class__.__name__print(type(model).__name__) # MyNet这在保存多个模型、按类名路由到不同加载逻辑时特别有用。对排查问题也很有帮助当模型被层层包装后打印model可以快速确认它到底是不是预期的那个结构。3.5 state_dict 与保存加载模型不等于权重训练好的模型需要保存。新手最容易犯的错误是直接用torch.save(model, model.pt)保存整个对象。这样做在短时间内能加载回来但一旦你的代码结构变了、类路径变了加载就会失败毕竟 Python 对象序列化带有很强的运行时耦合。更稳妥的方式是保存state_dicttorch.save(model.state_dict(), model_weights.pth) # 加载时先创建相同结构的模型再读权重 model MyNet() model.load_state_dict(torch.load(model_weights.pth, weights_onlyTrue))state_dict里保存的是每个参数的名称和数值不包含代码结构所以跨环境、跨同事的代码版本都更安全。更进一步如果你要保存完整的训练状态包括 epoch、优化器状态、学习率调度器状态应该存成一个字典torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), }, checkpoint.pth)这是训练中断后恢复现场的规范姿势。4. 训练闭环数据集、损失函数和优化器是怎么串起来的模型结构写好了接下来就是训练。训练不是一个独立的“黑盒操作”而是由数据加载、损失计算、反向传播、参数更新组成的循环。很多人把注意力全放在模型结构上忽略了数据流水线结果模型没跑几步就因为数据格式问题报错。这里我完整走一遍训练闭环用 MNIST 作为例子让整套流程跑起来。4.1 自定义 Dataset 的三个关键点PyTorch 的Dataset是一个抽象类只需实现两个方法__len__和__getitem__。但实际使用中有三个关键点值得留意。第一__getitem__返回的数据通常是原始样本和标签你可以在这里做数据增强、归一化也可以不做等DataLoader里再处理。第二返回的数据尽量转成torch.Tensor因为DataLoader需要把多个样本堆叠成 batch如果返回的是 Python 列表、不同长度的字符串后续会很麻烦。第三__getitem__里的逻辑尽量轻量不要在裁剪图像时每次都重新读一遍磁盘大文件不然num_workers再高也扛不住。一个实用示例from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, df, transformNone): self.df df self.transform transform def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] x row[[feat1, feat2, feat3]].values.astype(float32) y int(row[label]) x torch.tensor(x) y torch.tensor(y) if self.transform: x self.transform(x) return x, y4.2 DataLoader 参数和 collate_fn 的故事DataLoader把Dataset输出的单个样本打包成 batch。最常用的几个参数batch_size决定每个 batch 多少样本shuffleTrue在每个 epoch 开始时打乱数据num_workers决定子进程数pin_memoryTrue在 GPU 训练时可以让数据拷贝更快drop_lastTrue丢弃最后不足一个 batch 的样本避免某些 BatchNorm 层在 batch size1 时崩溃。遇到变长序列时光靠__getitem__返回不同长度的张量是不行的DataLoader默认的collate_fn会尝试用torch.stack把它们堆在一起结果直接报错。此时要自定义collate_fn用pad_sequence补齐from torch.nn.utils.rnn import pad_sequence def collate_fn(batch): xs, ys zip(*batch) xs_padded pad_sequence(xs, batch_firstTrue, padding_value0) ys torch.stack(ys) return xs_padded, ys4.3 一个能跑的 MNIST CNN 实例从零跑通第一个 epoch用 MNIST 搭一个简单 CNN 来跑完整训练循环。这里我把所有关键环节都放一起方便对照。import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms device torch.device(cuda if torch.cuda.is_available() else cpu) transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_ds datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) val_ds datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_ds, batch_size64, shuffleTrue, num_workers2) val_loader DataLoader(val_ds, batch_size512, shuffleFalse, num_workers2) class Net(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, kernel_size3) self.conv2 nn.Conv2d(32, 64, kernel_size3) self.fc1 nn.Linear(64 * 5 * 5, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x F.relu(self.conv1(x)) x F.max_pool2d(x, 2) x F.relu(self.conv2(x)) x F.max_pool2d(x, 2) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) return self.fc2(x) model Net().to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(3): model.train() total_loss 0.0 for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() outputs model(x) loss criterion(outputs, y) loss.backward() optimizer.step() total_loss loss.item() model.eval() correct 0 total 0 with torch.no_grad(): for x, y in val_loader: x, y x.to(device), y.to(device) outputs model(x) preds outputs.argmax(dim1) correct (preds y).sum().item() total y.size(0) print(fepoch {epoch 1}, loss: {total_loss / len(train_loader):.4f}, acc: {correct / total:.4f})这里每个细节都有为什么。model.train()让 BatchNorm/Dropout 进入训练模式model.eval()让它们进入推理模式torch.no_grad()关闭计算图outputs.argmax(dim1)取概率最大的类别下标。这套循环可以原封不动搬到自己的数据集上只需要换掉 Dataset。4.4 新手在训练循环里最容易踩的三个隐形坑第一个坑忘写optimizer.zero_grad()。梯度默认累积每多跑一个 batch梯度就在旧梯度上叠加。表现出来就是 loss 不降反升或者出现一个 batch 正常、另一个 batch 爆炸的周期性抖动。第二个坑model.eval()忘记调用。只要模型里有 Dropout 或 BatchNorm推理时不开 evalDropout 还在随机丢弃、BatchNorm 还在用 batch 统计量推理结果会不稳定。很多人验证集准确率忽高忽低往往就是这个原因。第三个坑把.item()到处乱用。有些初学者为了看中间变量把所有张量都转成 Python 数字结果断开了计算图。想监控某个中间激活可以保留张量在no_grad下打印它的mean()之类的标量值没必要全局.item()。5. 从常见代码到生产落地Attention、LSTM 源码与 ONNX 导出模型搭建不只是写 MVP 训练脚本还包括扩展到序列模型和最终部署。我会把三个高频需求放在一起讲手写通用 Attention 模块、理解 LSTM 源码到底在干什么、以及把训练好的模型导出成 ONNX。5.1 一个通用 Seq2Seq Attention 模块的手写与验证Attention 现在的应用很广但很多图森破的教程把整个注意力机制包装成了一个黑盒。实际上从 PyTorch 的角度看它就是一个普通nn.Module输入编码器的输出和当前解码器状态输出一个“加权上下文向量”和注意力权重。下面是一个加性注意力Additive Attention / Bahdanau Attention的简化实现class Attention(nn.Module): def __init__(self, enc_hidden, dec_hidden, attn_dim): super().__init__() self.W_e nn.Linear(enc_hidden, attn_dim, biasFalse) self.W_d nn.Linear(dec_hidden, attn_dim, biasFalse) self.v nn.Parameter(torch.randn(attn_dim)) def forward(self, enc_outs, dec_state): # enc_outs: [src_len, batch, enc_hidden] # dec_state: [batch, dec_hidden] pe self.W_e(enc_outs) # [src_len, batch, attn_dim] pd self.W_d(dec_state).unsqueeze(0) # [1, batch, attn_dim] scores torch.tanh(pe pd).matmul(self.v) # [src_len, batch] weights torch.softmax(scores, dim0) # [src_len, batch] context (enc_outs * weights.unsqueeze(-1)).sum(dim0) # [batch, enc_hidden] return context, weights逐行解释一下编码器所有时间步的输出首先被线性映射到一个维度为attn_dim的空间解码器当前隐状态也映射到同一空间二者相加后过tanh再与向量v做点积得到每个编码器位置的“分数”softmax得到权重最后用权重对所有编码器输出加权求和得到上下文向量。这段代码在原版 Seq2Seq 里通常是这样用的解码器每个时间步计算 attention然后把上下文向量与当前解码器输入拼接再送入 RNN 单元。你也可以把这段代码搬去 Transformer 里的 cross-attention只是要把加性注意力换成缩放点积注意力。核心思维一致Attention 就是一个可微分的查询函数输入 query 和 key-value 对输出加权后的 value。5.2 读 PyTorch LSTM 源码时你应该关注哪几个点很多人在网上搜pytorch lstm源码但打开 GitHub 后发现底层是 C/CUDA 的融合实现看得一脸懵。我的建议是读nn.LSTM不一定要逐行读底层而是先搞清楚它的接口行为和返回结构。lstm nn.LSTM(input_size10, hidden_size20, num_layers2, batch_firstTrue) x torch.randn(5, 8, 10) # batch5, seq_len8, input_size10 out, (h_n, c_n) lstm(x)这里的返回结构很重要out是所有时间步的最后一层输出形状[batch, seq_len, hidden_size]batch_firstTrue时h_n是最后一个时间步每个层的隐状态形状[num_layers, batch, hidden_size]c_n是最后一个时间步每个层的细胞状态形状同上。如果你好奇双向 LSTM输出维度会翻倍成hidden_size * 2。读源码时最值得关注的是forward里如何处理输入长度、如何初始化隐状态、以及它怎么调用_VF.lstm这个融合算子。对大多数人而言把nn.LSTM当nn.Module来使用、理解形状变换就够了。真正想深入内部再去看手动循环实现nn.LSTMCell的方式。还有一个实际经验不要轻易在forward里用 LSTM 的输出直接接全连接层时忽略out[:, -1, :]和h_n[-1]的区别。前者是最后一个时间步的隐含输出后者是最后一个层、最后一个时间步的隐状态。对于单向 LSTM两者通常是一致的但在多层或双向 LSTM 中取法就完全不同。5.3 导出 ONNX模型离开 PyTorch 环境的最后一公里模型训练完常常要部署到服务端推理框架比如 ONNX Runtime、TensorRT。PyTorch 提供了torch.onnx.export但它并不只是写一行代码那么简单。最常用的导出模板长这样model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], dynamic_axes{ input: {0: batch}, output: {0: batch}, }, )几个关键点模型必须切到eval()模式否则导出的图会保留训练时的不确定性行为dummy_input的尺寸和迁移设备必须与实际推理一致opset_version决定了导出算子集合想兼容较老的推理版本就别一味追新dynamic_axes用于把 batch 维设成动态否则导出后只能接受固定 batch size。导出后务必用onnxruntime验证一次比较推理结果与 PyTorch 输出的误差import onnxruntime as ort import numpy as np ort_sess ort.InferenceSession(model.onnx) x_cpu dummy_input.cpu().numpy() outputs ort_sess.run(None, {input: x_cpu})[0]如果误差在 1e-5 数量级说明导出成功。误差大往往与 BatchNorm、Dropout 模式或自定义算子有关。6. 设备适配与运行兼容实战解决“torch 不支持设备”和一堆环境怪问题最后这部分集中讲运行阶段最让人头疼的设备适配问题。我亲眼见过有人因为硬件或环境的问题折腾一个礼拜最后发现只是装错了版本。6.1 没有 GPU 能不能学 PyTorch很多人把加速当成了必需先给结论能而且非常建议。没有 NVIDIA GPU装 CPU 版 PyTorch 一样能学完模型搭建、训练循环、反向传播原理这些全部核心内容。区别只是大模型训练慢一些、batch size 小一些。对入门来说把 MNIST、CIFAR 这类小数据集跑通CPU 完全够用。Apple Silicon 用户还可以通过 MPS 后端把计算放到 GPU 上torch.backends.mps.is_available()为 True 时把 device 设成mps即可。AMD 用户则走 ROCm 路线后面细说。6.2 “不支持设备”类报错的排查链路以绘世启动器为例很多玩绘图整合包的读者会遇到绘世启动器显示“PyTorch 不支持设备”之类的问题。这通常不是整合包坏了而是它继承的 PyTorch 版本是 CPU 版或者 CUDA 算力低于 PyTorch 要求。排查链路我建议这样走第一步在整合包的环境里跑这一段import torch print(torch.__version__) print(torch.version.cuda) print(torch.cuda.is_available())第二步如果torch.cuda.is_available()是 False先跑nvidia-smi确认驱动能看到显卡。如果驱动能看到但 PyTorch 看不到多半是 PyTorch 装成了 CPU 版或者显卡算力太低。比如 Maxwell 之前的架构新版本 PyTorch 已经不再支持需要换低版本 torch 或用 CPU 推理。第三步如果确实是 CPU 版就用对应 CUDA 版本重新安装pip uninstall torch torchvision torchaudio pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121有些启动器里有“版本检测”逻辑它会扫描环境变量里的 Python 和 torch也可能会因为虚拟环境路径不对而误判。如果你确定 torch 装对了但启动器仍然报错检查启动器使用的是不是同一个 conda 环境这是最常见的人为乌龙。6.3 7900XTX 加 WSLAMD GPU 跑 torch 的真实体验最近 AMD 用户在社区问 7900XTX 跑 PyTorch 怎么弄我实际试过的路径是 WSL2 ROCm 版 PyTorch。AMD 显卡不能直接用 NVIDIA 的 CUDA 版本需要安装 ROCm 构建的 torch。安装方式pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6在 PyTorch 的 ROCm 命名空间里代码层面依然使用torch.cuda.is_available()来判断设备这个返回 True 但底层实际走的是 ROCm/HIP。WSL2 里跑的一个常见注意事项是显卡驱动要装在 Windows 侧WSL 内部不需要重复安装驱动只需确认 WSL 内核和 Windows 版本够新。另外WSL2 默认内存可能偏小大模型容易在申请显存时先触发系统 OOM 报错可以考虑在.wslconfig里给足内存。真实体感配置好以后跑 PyTorch 和 NVIDIA 卡的体验几乎一致但第三方优化算子适配目前还是不如 CUDA 生态丰富编译原生存量算子时偶尔会缺头文件需要多留点余量。6.4 CentOS7 上通过 Anaconda 装 torch 的注意事项有段时间我需要在 CentOS7 机器上配置 PyTorch 环境步骤本身不复杂先装 Anaconda再创建 conda 环境再装 torch。真正的坑在老系统的系统库兼容性上。CentOS7 自带的glibc版本偏低PyTorch 官方 wheel 对glibc版本有最低要求。比较新的 torch 2.x 版本在这些老系统上可能装不上或者装上后 import 时报GLIBC_2.27 not found。这时候优先选低版本 torch比如 torch 1.8/1.9/1.10或者换一台系统库较新的机器。另外一个更稳的路径是在 CentOS7 上通过 conda 安装因为 conda 打包的依赖相对完整会减少一部分系统库冲突问题。安装 Anaconda 的命令是通用的wget https://repo.anaconda.com/archive/Anaconda3-2023.09-0-Linux-x86_64.sh bash Anaconda3-2023.09-0-Linux-x86_64.sh source ~/.bashrc conda create -n torch python3.10 conda activate torch如果 Anaconda 源下载慢可以换成国内镜像源再装。这类老环境最值得记住的教训是先确认操作系统基础库版本再挑 torch 版本别一上来就用最新版。我在 CentOS7 上踩过的坑绝大多数是基础库版本跟不上 torch 新特性。6.5 升级 PyTorch 后常见的兼容性问题最后说一个高频现象把 torch 从 1.x 升到 2.x 后原有代码突然报错或行为变化。原因可能是某些旧 API 被移除、默认行为改变或者第三方库版本还没跟上。常见的几个点torch.load默认的weights_only行为变化某些老权重文件加载时需要显式设置weights_onlyFalsetorch.range被废弃改成torch.arange某些自定义算子在编译时依赖旧版 C ABI升级后需要重新编译nn.Module的.to(cuda)行为没有变化但torch.cuda.set_device的使用方式建议改成device参数统一控制。我的建议是读升级日志看 release note 中的 breaking change永远比靠猜靠谱。官方文档里专门有一节讲升级兼容性升级前翻一翻能省掉大量 debug 时间。最后再讲一个个人经验。我后来做任何新项目开头固定会花十分钟做环境和结构验证建独立 conda 环境、打印 torch 版本和设备、写一个最小的两层网络打印state_dict的键和形状。这套“最小验证流程”跑通了再往里面加数据流水线和复杂模型。很多人在模型结构里改了半天最后发现问题是整套环境从起点就不对那才是最浪费时间的。你如果现在正被某个环境或设备报错卡住不妨先退回这个最小流程用排除法确认 torch 本身能正常工作再继续往下走。环境这层账早算清楚后面省下的时间真的非常多。