3. 模型序列化：从PyTorch/TensorFlow导出ONNX，ONNX转TensorRT引擎

好，咱们进入第三章。这一章可以说是整个TensorRT优化流程的「咽喉要道」。模型训练得再好，导不出来、转不过去，后面全是白搭。我见过太多团队在训练上花了几个月，最后卡在模型转换这一步，一卡就是一两周。

说白了，序列化就是把你的PyTorch或TensorFlow模型，变成TensorRT能吃的「饭」。这顿饭怎么做？分两步：先导出ONNX，再把ONNX转成TensorRT引擎。

3.1 为什么非要ONNX？

你可能会问：TensorRT不能直接吃PyTorch模型吗？嗯，严格来说可以，通过torch2trt或者直接解析TorchScript。但我个人习惯，还是走ONNX这条路。为什么？

• 通用性：ONNX是模型界的「普通话」。不管你是PyTorch、TensorFlow还是其他框架，都能转成ONNX。以后换框架，模型不用重写。
• 调试方便：ONNX有可视化工具（Netron），你可以直接看到计算图长什么样。我在项目中遇到过好几次，模型精度不对，一查ONNX图，发现某个算子被合并错了。
• 算子覆盖全：TensorRT对ONNX算子的支持最成熟。你走TorchScript路线，有些自定义算子可能就卡住了。

3.2 从PyTorch导出ONNX

PyTorch导出ONNX，主要靠 torch.onnx.export。这个函数看起来简单，但坑不少。我直接上代码，然后讲关键点。

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = MyModel().eval()  # 一定要eval模式！
dummy_input = torch.randn(1, 3, 224, 224)  # 一个batch的输入

torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    "model.onnx",             # 输出文件名
    export_params=True,       # 导出参数
    opset_version=17,         # 算子集版本，我建议用17或18
    do_constant_folding=True, # 常量折叠，能省则省
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},   # 动态batch
        'output': {0: 'batch_size'}
    }
)

这里有几个点，我得重点说一下：

3.2.1 opset_version 选哪个？

我个人习惯用17或18。为什么？因为TensorRT 8.6及以上版本对opset 17支持得最好。你选太低的版本，有些新算子不支持；选太高，TensorRT可能还没跟上。我踩过这个坑：有一次用了opset 20，结果TensorRT报错说某个算子不认识，折腾了半天才找到原因。

3.2.2 dynamic_axes 一定要设

你想想看，生产环境里，你的输入batch size会固定吗？不会。所以一定要设dynamic_axes。这样导出的ONNX支持动态batch。不过要注意，TensorRT在构建引擎时，需要指定一个范围，比如min=1, opt=4, max=8。这个后面会讲。

3.2.3 检查ONNX图

导出之后，别急着下一步。先用工具检查一下：

import onnx
from onnx import checker

model = onnx.load("model.onnx")
checker.check_model(model)  # 检查模型结构
print(onnx.helper.printable_graph(model.graph))  # 打印计算图

我建议你用Netron（一个网页工具）打开ONNX文件，看看图是不是你想要的。有一次我发现模型里多了一个没用的Identity节点，就是靠Netron看出来的。

3.3 从TensorFlow导出ONNX

TensorFlow这边，主要用tf2onnx工具。流程类似，但有些细节不同。

import tf2onnx
import tensorflow as tf

# 加载你的TensorFlow模型
model = tf.saved_model.load("path/to/saved_model")

# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 导出ONNX
output_path = "model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    opset=17,
    output_path=output_path
)

这里要注意：TensorFlow的模型保存方式有好几种（SavedModel、H5、Checkpoint）。我建议统一用SavedModel格式，因为tf2onnx对它的支持最稳定。我曾经用H5格式导出，结果某些层被解析错了，换成SavedModel就好了。

3.4 ONNX转TensorRT引擎

ONNX到手了，接下来就是重头戏：转成TensorRT引擎。这一步用trtexec工具或者TensorRT Python API都可以。我一般用Python API，因为可以精细控制。

3.4.1 用Python API转换

import tensorrt as trt

# 创建logger和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)

# 解析ONNX模型
with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise RuntimeError("ONNX解析失败")

# 创建配置
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 设置动态形状
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 224, 224), (4, 3, 224, 224), (8, 3, 224, 224))
config.add_optimization_profile(profile)

# 构建引擎
serialized_engine = builder.build_serialized_network(network, config)

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

这段代码里，有几个关键点：

• EXPLICIT_BATCH：必须加这个flag，否则不支持动态batch。我刚开始做的时候忘了加，结果引擎只能跑固定batch，被坑惨了。
• 优化profile：这里设置了三个点：min、opt、max。TensorRT会根据这些点做层融合和内存优化。opt一般设为你实际部署时最常用的batch size。
• workspace大小：1GB是常见设置。如果你的模型很大，可以调高。但注意，不是越大越好，太大了可能反而导致内存碎片。

3.4.2 用trtexec命令行工具

如果你不想写代码，也可以用trtexec。它随TensorRT一起安装。

trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:4x3x224x224 \
        --maxShapes=input:8x3x224x224 \
        --workspace=1024 \
        --fp16

这个命令和上面的Python API效果一样。加--fp16可以启用FP16推理，速度能翻倍。不过要注意，有些模型在FP16下精度会掉，需要验证。

3.5 验证引擎

引擎生成之后，一定要验证。我一般做两件事：

1. 精度验证：用同一批输入，对比PyTorch模型和TensorRT引擎的输出。误差在1e-3以内算正常。
2. 性能测试：用trtexec跑一下benchmark，看看吞吐量和延迟。

trtexec --loadEngine=model.engine --shapes=input:4x3x224x224 --duration=10

这个命令会跑10秒，输出平均延迟和吞吐量。我一般看两个指标：

指标	说明	参考值
Latency (ms)	单次推理延迟	< 10ms（实时场景）
Throughput (qps)	每秒处理请求数	越高越好

如果延迟太高，可以考虑开启FP16或者INT8量化。不过那是后面章节的内容了。

3.6 避坑指南

最后，我总结几个常见坑，都是我曾经踩过的：

• 模型没设eval模式：PyTorch模型导出前一定要model.eval()，否则BN和Dropout的行为会乱掉。
• 输入形状不匹配：ONNX里的输入形状和TensorRT profile里的形状必须一致。比如ONNX里是NCHW，profile里也得是NCHW。
• 算子版本冲突：ONNX的opset和TensorRT支持的opset要匹配。我建议用opset 17，兼容性最好。
• 动态形状范围太大：min和max的范围不要太大，否则TensorRT会生成很多优化方案，导致构建时间变长。一般min=1, max=8就够了。

嗯，这一章的内容就到这里。模型序列化是基础，但也是最容易出问题的一步。你只要把ONNX导干净，TensorRT引擎的构建就成功了一大半。下一章，我们会讲引擎的加载和推理，到时候这些知识都会用上。