模型导出与序列化：PyTorch模型保存与加载、ONNX导出、TorchScript追踪与脚本化、模型格式转换

模型训练完了，只是第一步。真正要把它扔到生产环境里跑，还得过「序列化」这一关。说白了，就是把训练好的模型参数和结构，打包成一种标准格式，让别的平台也能读懂、跑起来。

我个人习惯把这一章叫做「模型出厂前的最后一道工序」。你想想看，你在PyTorch里调参调得再好，到了C++推理引擎、移动端或者嵌入式设备上，PyTorch环境都没了，怎么跑？所以，导出和序列化，就是给模型换上一套「通用制服」。

PyTorch模型保存与加载：最基础的两种姿势

PyTorch官方提供了两种保存方式，我刚开始学的时候也迷糊过一阵子。咱们直接看代码。

方式一：保存/加载整个模型（不推荐）

# 保存
torch.save(model, 'model.pth')

# 加载
model = torch.load('model.pth')

这种方式简单粗暴，但有个大坑——它把模型的定义和参数打包在一起了。如果你换了PyTorch版本，或者模型类定义改了，加载时直接报错。我在项目里吃过这个亏，升级了PyTorch小版本，结果之前保存的模型全废了。

方式二：保存/加载状态字典（推荐）

# 保存
torch.save(model.state_dict(), 'model_state_dict.pth')

# 加载
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load('model_state_dict.pth'))
model.eval()

这才是正经做法。只保存参数，不保存结构。加载时你需要先实例化模型，再把参数灌进去。嗯，这里要注意：加载完一定要调 model.eval()，否则Dropout和BatchNorm的行为会乱掉。我曾经有一次忘了调，线上推理结果忽高忽低，排查了半天才发现是这行代码漏了。

ONNX导出：让模型走出PyTorch的围墙

ONNX（Open Neural Network Exchange）是微软和Facebook联手搞的中间格式。说白了，它就像模型界的「普通话」——不管你是PyTorch、TensorFlow还是其他框架，都能转成ONNX，然后被各种推理引擎消费。

为什么要导出ONNX？我个人觉得最大的好处是：推理加速。ONNX Runtime、TensorRT这些引擎，对ONNX的支持非常成熟，推理速度能快好几倍。

导出ONNX的标准流程

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = YourModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

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

这里有个关键点：dynamic_axes。如果你希望模型支持动态batch大小，一定要加上这个参数。我刚开始做导出时没加，结果线上请求的batch大小一变，ONNX Runtime直接报错。

TorchScript：PyTorch自家的序列化方案

ONNX虽然通用，但有时候PyTorch特有的算子ONNX不支持，导出会报错。这时候TorchScript就派上用场了。它是PyTorch自家的序列化格式，能完整保留模型的计算图。

TorchScript有两种生成方式：追踪（Tracing）和脚本化（Scripting）。

追踪（Tracing）

# 用虚拟输入跑一遍，记录计算图
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save('model_traced.pt')

追踪方式很简单，但有个致命弱点：它只能记录实际执行过的路径。如果模型里有条件分支，没走到的分支就不会被记录。我有个模型在推理时根据输入长度选择不同处理逻辑，用追踪方式导出后，线上遇到没见过的输入长度直接崩了。

脚本化（Scripting）

# 用脚本化方式，编译器会解析整个模型代码
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

脚本化方式更安全，它会解析整个模型的Python代码，生成完整的计算图。代价是有些Python语法不支持，比如动态的字典操作、某些内置函数。嗯，这里需要一点学习成本。

模型格式转换：.pt → .onnx → .trt

这是生产环境中最常见的转换链路。PyTorch模型先转成ONNX，再用TensorRT优化成TensorRT引擎（.trt），部署到NVIDIA GPU上跑。

为什么要多这一步？TensorRT会对计算图做融合、量化、内存优化，推理速度能再提升2-5倍。我在一个目标检测项目里，从PyTorch直接推理到TensorRT优化后，延迟从30ms降到了8ms，效果非常明显。

ONNX转TensorRT

import tensorrt as trt

# 创建TensorRT builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 解析ONNX模型
parser = trt.OnnxParser(network, logger)
with open('model.onnx', 'rb') as f:
    parser.parse(f.read())

# 构建引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
serialized_engine = builder.build_serialized_network(network, config)

# 保存引擎
with open('model.trt', 'wb') as f:
    f.write(serialized_engine)

知识体系总览

下面这张图，把整个模型导出与序列化的流程串起来了。你可以把它当作一张「路线图」来用。

总结一下

模型导出与序列化，说白了就是给模型办「护照」。PyTorch的 .pt 是身份证，只能在PyTorch环境里用；ONNX是国际护照，大部分推理引擎都认；TensorRT的 .trt 是VIP通行证，专门给NVIDIA GPU加速用。

我个人建议的实践顺序是：先用 state_dict 保存PyTorch模型，然后尝试导出ONNX。如果ONNX报错，退一步用TorchScript脚本化。最后，如果部署在NVIDIA GPU上，再转成TensorRT引擎。每一步都有坑，但踩过一次就记住了。

---