豆包AI怎么写PyTorch模型｜代码演示教程

PyTorch模型保存与加载远非简单的“存一下、读回来”——看似一行代码就能搞定的torch.save和torch.load，实则暗藏设备不匹配、结构错位、计算图残留、trace失真等多重陷阱：用state_dict()才是轻量安全的常规选择，而直接保存整个model极易导致eval失效、梯度冲突或跨环境加载失败；加载时遇到Missing keys或设备错误，根源常在类定义变更、DataParallel前缀、GPU权重未映射到CPU；至于JIT导出结果异常，往往源于trace固化了特定输入下的控制流。真正可靠的模型IO，靠的不是侥幸跑通，而是每一步都直面PyTorch底层序列化、设备抽象与图构建机制的精密博弈。

PyTorch模型保存时用 torch.save(model.state_dict(), ...) 还是 torch.save(model, ...)

用 state_dict() 是常规且安全的选择，直接保存整个 model 容易出问题。

常见错误现象：加载用 torch.save(model, path) 保存的模型后，model.eval() 不生效、Dropout 仍在训练模式，或反向传播报 RuntimeError: Trying to backward through the graph a second time —— 因为保存了计算图和缓存的梯度状态。

• state_dict() 只存参数张量，轻量、可移植、不绑定模型类定义，适合部署和复训
• 直接保存 model 会序列化整个 Python 对象，依赖当前代码结构；换一个文件里类名/路径稍有不同就 ModuleNotFoundError 或 AttributeError
• 如果必须保存完整模型（比如含自定义 forward 逻辑且不想重建实例），务必确保加载时环境完全一致，并手动调用 model.train(False) 清理状态

加载模型时 load_state_dict() 报 Missing key(s) in state_dict 怎么办

说明模型结构和保存时的结构对不上，不是路径错了，是定义变了。

典型场景：你在原模型上加了新层、改了层名（比如把 self.fc 改成 self.classifier），但没更新保存逻辑；或者用了 nn.DataParallel 训练后，在单卡上直接加载 —— 此时键名会多出 module. 前缀。

• 先用 print(list(model.state_dict().keys())) 和 print(list(checkpoint.keys())) 对比键名差异
• 如果是 module. 前缀问题，加载前做一次映射：{k.replace("module.", ""): v for k, v in checkpoint.items()}
• 如果只是多了几个键（比如新增的 self.aux_head），加 strict=False 参数跳过不匹配项，但得自己确认哪些没载入
• 别用 model.load_state_dict(checkpoint, strict=True) 硬扛，报错就停，它不会帮你猜意图

在 CPU 上加载 GPU 训练的模型，为什么报 Expected all tensors to be on the same device

因为 torch.load() 默认把权重放在原设备上，GPU 保存的 checkpoint 里张量还是 cuda:0 类型，CPU 模型没法直接接。

这不是模型写错了，是加载姿势不对。哪怕你后续调用 model.cpu()，也得先让 load_state_dict() 成功执行完。

• 加载时指定 map_location：用 torch.load(path, map_location="cpu") 或 map_location=torch.device("cpu")
• 如果模型是在多卡上用 DistributedDataParallel 训练的，保存的是 model.module.state_dict()，加载时也要注意别漏掉 .module
• 别等 forward() 跑起来才报错，这个错误一定发生在 load_state_dict() 阶段之后、第一次 model(input) 之前

用 torch.jit.trace 导出模型后，为什么推理结果和原始模型不一致

trace 会固化模型某次前向过程中的控制流和张量形状，一旦输入变了，行为就不可靠。

常见于含条件分支（if x.size(0) > 1:）、动态 padding、或依赖输入值做 shape 推断的操作（比如 x.view(x.size(0), -1) 在 trace 时若 batch=1，-1 就被算死成某个固定数）。

• trace 前确保输入是典型尺寸 + 典型值，比如 batch=2、带真实文本/图像内容，别用全零张量
• 避免在 forward 里写 Python 控制流；该用 torch.where、torch.nn.functional.pad 等可 trace 操作替代
• 导出后务必用相同输入跑一遍原始模型和 traced 模型，用 torch.allclose(out1, out2, atol=1e-5) 校验，别只看 shape
• 如果模型逻辑复杂，优先考虑 torch.jit.script，但它要求代码能被静态分析，对 typing 和 control flow 更敏感

今天关于《豆包AI怎么写PyTorch模型｜代码演示教程》的内容介绍就到此结束，如果有什么疑问或者建议，可以在golang学习网公众号下多多回复交流；文中若有不正之处，也希望回复留言以告知！