模型导出与格式转换：从PyTorch到TorchScript、ONNX再到TensorRT

模型训练好了，只是万里长征走了一半。真正要上线跑起来，还得过「格式转换」这一关。我刚开始做部署那会儿，就吃过格式不兼容的亏——模型在GPU上跑得飞快，换到推理引擎上直接报错。嗯，今天咱们就把这块硬骨头啃下来。

为什么需要模型导出？

说白了，PyTorch训练出来的模型是个「富家子弟」——依赖一大堆Python环境、动态图机制、还有各种调试信息。但生产环境要的是「轻装上阵」：

• 去掉Python依赖：C++部署不需要装PyTorch
• 静态化计算图：推理时不需要动态构建图，速度更快
• 量化与优化：FP16、INT8这些压缩手段，原生PyTorch支持有限

我个人习惯把模型导出分成三步走：TorchScript → ONNX → TensorRT。每一步都有它的用武之地，别想着一步到位。

TorchScript：PyTorch的亲儿子

TorchScript是PyTorch官方推出的中间表示。它能把动态图「凝固」成静态图，同时保留大部分PyTorch的操作语义。

两种导出方式

方式一：Tracing（追踪）

给模型一个样例输入，PyTorch会记录下所有执行过的操作。简单粗暴，但遇到控制流（if/for）就抓瞎。

import torch

class MyModel(torch.nn.Module):
    def forward(self, x):
        return x * 2 + 1

model = MyModel()
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

方式二：Scripting（脚本化）

直接解析Python代码，生成静态图。能处理控制流，但要求代码必须是「TorchScript兼容」的——不能用numpy、不能用Python内置的list操作。

@torch.jit.script
def my_func(x):
    if x.sum() > 0:
        return x * 2
    else:
        return x + 1

scripted_model = torch.jit.script(MyModel())
scripted_model.save("model_scripted.pt")

ONNX：模型界的通用语言

ONNX（Open Neural Network Exchange）是微软和Facebook联合推出的开放格式。它的目标很明确——让模型能在不同框架间自由迁移。

PyTorch导出ONNX

import torch.onnx

model = MyModel()
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=11
)

这里有个关键参数——dynamic_axes。它告诉ONNX：我这个模型的batch_size是可变的。如果不设置，导出的模型只能接受固定batch size的输入。

ONNX的常见坑

• 算子不支持：PyTorch里有些操作（如torch.einsum）ONNX没有对应实现。解决办法是用torch.onnx.symbolic自定义映射。
• 动态shape问题：ONNX对动态shape的支持有限。如果模型里有reshape操作，导出时最好固定住shape。
• 调试工具：用onnxruntime验证导出的ONNX模型是否正确。我习惯写个脚本，对比PyTorch和ONNX的输出差异。

import onnxruntime as ort
import numpy as np

# 验证ONNX模型
ort_session = ort.InferenceSession("model.onnx")
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = ort_session.run(None, {"input": input_data})
print("ONNX输出:", outputs[0])

TensorRT：NVIDIA的推理加速器

TensorRT是NVIDIA推出的推理优化引擎。它能把模型「榨干」到极致——FP16推理、INT8量化、层融合、内存优化……说白了，就是让GPU跑得更快。

从ONNX到TensorRT

最常用的方式是通过trtexec命令行工具：

trtexec --onnx=model.onnx \
        --saveEngine=model.trt \
        --fp16 \
        --workspace=4096 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:16x3x224x224

参数说明：

• --fp16：启用半精度推理，速度提升约2倍
• --workspace：指定工作空间大小（MB）
• --minShapes/optShapes/maxShapes：定义动态batch的范围

Python API方式

如果你需要在代码中动态构建TensorRT引擎，可以用Python API：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_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, 4 << 30)  # 4GB
config.set_flag(trt.BuilderFlag.FP16)

serialized_engine = builder.build_serialized_network(network, config)
with open("model.trt", "wb") as f:
    f.write(serialized_engine)

TensorRT的量化

FP16是入门，INT8才是真正的性能利器。但INT8量化需要校准数据集：

# INT8量化配置
config.set_flag(trt.BuilderFlag.INT8)
calibrator = MyCalibrator(calibration_data)
config.int8_calibrator = calibrator

校准数据集一般选500-1000张有代表性的图片。选得不好，精度会掉得很厉害。我遇到过最夸张的一次，INT8量化后精度掉了5个点，后来换了校准集才救回来。

三种格式的对比

特性	TorchScript	ONNX	TensorRT
依赖	PyTorch C++库	ONNX Runtime	TensorRT
推理速度	中等	中等	最快
量化支持	有限	一般	FP16/INT8
动态shape	支持	有限	支持
调试难度	低	中	高

我的推荐流程

在实际项目中，我一般这样选：

1. 快速验证：用TorchScript，省事
2. 跨平台部署：导出ONNX，用ONNX Runtime跑
3. 极致性能：ONNX转TensorRT，上FP16或INT8

你想想看，如果模型只在NVIDIA GPU上跑，直接上TensorRT就完事了。但如果要部署到CPU、AMD GPU或者手机端，ONNX才是王道。