3、模型导出基础：PyTorch模型保存与加载、ONNX导出原理、TorchScript简介

好，咱们进入第三章。这一章我把它叫做「模型部署的起跑线」。你想想看，模型训练得再好，如果导不出来、导出来跑不动，那一切都是白搭。我在项目里见过太多人，训练时各种花活，一到导出就卡壳——嗯，今天咱们就把这块硬骨头啃下来。

3.1 PyTorch模型保存与加载：不只是torch.save

先说最基础的。PyTorch保存模型，说白了就两种方式：保存整个模型，或者只保存状态字典。我个人习惯用后者，为什么？因为灵活。

3.1.1 两种保存方式对比

我曾经在一个客户现场遇到过这么个事：对方把整个模型用 torch.save 存了，结果换了一台机器，模型类定义路径变了，直接报错。折腾了两个小时才找到原因。从那以后，我团队里谁再用第一种方式保存，我就请他喝咖啡——然后当面改代码。

3.1.2 加载时的注意事项

加载模型时，有个细节很多人会忽略：设备映射。你训练时用 GPU，加载时如果机器没 GPU，直接报错。所以我的标准写法是这样的：

# 推荐写法：兼容 CPU/GPU
checkpoint = torch.load('model.pth', map_location='cpu')
model.load_state_dict(checkpoint)

# 如果你明确知道目标设备
checkpoint = torch.load('model.pth', map_location=device)

torch.save({
    'model_state_dict': model.state_dict(),
    'config': model_config,
    'input_shape': (1, 3, 224, 224),
    'normalize_params': mean_std,
}, 'checkpoint.pth')

3.2 ONNX导出原理：跨平台的「通用语言」

ONNX 是什么？说白了就是模型界的「普通话」。你训练用 PyTorch，对方推理用 TensorRT，中间怎么传？靠 ONNX。它定义了一套标准化的计算图表示，各家框架都能读。

3.2.1 导出核心流程

PyTorch 转 ONNX 就一个函数：torch.onnx.export()。但参数怎么配，学问很大。我直接上代码：

import torch
import torch.onnx

# 定义模型并加载权重
model = MyModel()
model.load_state_dict(torch.load('model.pth', map_location='cpu'))
model.eval()

# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    'model.onnx',             # 输出路径
    export_params=True,       # 导出参数
    opset_version=11,         # 算子集版本
    do_constant_folding=True, # 常量折叠优化
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里我重点说几个参数：

• opset_version：我建议用 11 或更高。版本太低，有些算子不支持；版本太高，老推理引擎可能不兼容。11 是个比较稳妥的选择。
• dynamic_axes：这个很关键。如果你希望导出的模型支持动态 batch 大小，必须配这个。我踩过坑——没配 dynamic_axes，结果线上请求 batch 大小变了，推理直接崩。
• do_constant_folding：建议打开。它会把一些固定计算提前算好，减少运行时开销。

3.2.2 导出后的验证

导出完别急着走，一定要验证。ONNX 官方提供了检查工具：

import onnx
import onnxruntime as ort

# 检查模型结构
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)

# 用 onnxruntime 跑一遍推理
ort_session = ort.InferenceSession('model.onnx')
outputs = ort_session.run(None, {'input': dummy_input.numpy()})
print('ONNX 推理结果:', outputs[0].shape)

3.3 TorchScript简介：PyTorch的亲儿子

TorchScript 是 PyTorch 自家的序列化格式。它不像 ONNX 那样追求跨平台，而是专注于 PyTorch 生态内的优化和部署。说白了，如果你整个链路都是 PyTorch，TorchScript 可能是更好的选择。

3.3.1 两种生成方式

TorchScript 有两种生成方式：trace 和 script。我分别说一下：

我个人习惯：能用 trace 就用 trace。为什么？因为 trace 简单、稳定，不容易出错。script 虽然功能强，但经常遇到「这个语法不支持」的尴尬。

3.3.2 trace 实战

# trace 方式
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save('model_traced.pt')

# 加载使用
loaded_model = torch.jit.load('model_traced.pt')
output = loaded_model(dummy_input)

你看，代码非常简洁。但有个坑：trace 只记录你给的那个输入路径。如果你的模型里有 if-else 分支，trace 只会记录走的那条路。我曾经有个模型，训练时走的是分支 A，导出时也走的 A，结果线上走了分支 B，直接输出乱码。嗯，这就是 trace 的局限性。

3.3.3 script 什么时候用

当你的模型包含以下情况时，必须用 script：

• 动态循环（比如根据输入长度决定循环次数）
• 条件分支（if-else 依赖输入值）
• 自定义的 Python 控制流

@torch.jit.script
def my_forward(x, threshold):
    if x.max() > threshold:
        return x * 2
    else:
        return x / 2

# 然后把这个函数集成到模型中

3.4 本章知识体系

为了让你更直观地理解这三者的关系，我画了一张图：

这张图把三种导出方式的关系理清楚了。你从训练模型出发，根据实际需求选择路径：要跨平台用 ONNX，要 PyTorch 生态优化用 TorchScript，要灵活加载用 state_dict。没有绝对的好坏，只有合不合适。

公众号:蓝海资料掘金营，微信deep3321