3、TensorRT模型导出基础：从PyTorch导出ONNX、ONNX模型检查与优化、TensorRT解析ONNX、Engine序列化与反序列化

好，咱们进入正题。这一章是整个课程的地基，说白了就是搞清楚模型怎么从PyTorch一路跑到TensorRT上。我见过太多同学在ONNX导出这一步就卡住了，后面全白搭。所以咱们慢慢来，把每个环节都掰扯清楚。

3.1 从PyTorch导出ONNX

PyTorch导出ONNX，核心就一个函数——torch.onnx.export()。但参数怎么填，坑在哪里，我踩过不少。

先看一个最基础的例子：

import torch
import torchvision.models as models

# 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

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

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    "resnet18.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'}
    }
)
print("ONNX导出完成")

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

• opset_version：我个人习惯用11或13。版本太低有些算子不支持，版本太高某些推理框架可能不兼容。11是个稳妥的选择。
• dynamic_axes：这个参数很关键。如果你希望模型支持动态batch，就得在这里指定。我在项目中遇到过，一开始没设这个，结果线上请求batch大小变化时，每次都得重新构建Engine，性能直接崩了。
• do_constant_folding：建议打开。它能提前把一些固定计算算好，减少ONNX模型里的冗余节点。

3.2 ONNX模型检查与优化

导出的ONNX文件，不能直接扔给TensorRT。你得先检查一下，看看有没有问题。我常用的工具是 onnxruntime 和 onnx 自带的检查器。

3.2.1 基础检查

import onnx

# 加载ONNX模型
onnx_model = onnx.load("resnet18.onnx")

# 检查模型结构是否合法
onnx.checker.check_model(onnx_model)
print("ONNX模型检查通过")

# 打印模型信息
print(onnx.helper.printable_graph(onnx_model.graph))

这一步能发现一些明显的问题，比如节点连接错误、类型不匹配等。但说实话，很多时候检查通过了，实际跑起来还是有问题。所以我会再用onnxruntime跑一遍推理，验证数值是否正确。

import onnxruntime as ort
import numpy as np

# 创建推理会话
ort_session = ort.InferenceSession("resnet18.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].shape)

3.2.2 模型优化

ONNX模型里经常有一些冗余操作，比如连续的Reshape、Cast等。我推荐用 onnx-simplifier 来清理：

pip install onnx-simplifier
python -m onnxsim resnet18.onnx resnet18_simplified.onnx

这个工具能合并一些节点，去掉不必要的操作。我在一个项目里，优化前模型有200多个节点，优化后只剩150个，推理速度提升了将近20%。

3.3 TensorRT解析ONNX

ONNX准备好了，接下来就是用TensorRT来解析它。TensorRT提供了一个叫 trtexec 的命令行工具，但咱们做工程化，还是用Python API更灵活。

import tensorrt as trt

# 创建logger
logger = trt.Logger(trt.Logger.WARNING)

# 创建builder
builder = trt.Builder(logger)

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

# 创建ONNX解析器
parser = trt.OnnxParser(network, logger)

# 解析ONNX模型
with open("resnet18_simplified.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解析失败")

print("ONNX解析成功，网络包含 %d 层" % network.num_layers)

这里有个细节：EXPLICIT_BATCH 这个标志。如果你在导出ONNX时用了dynamic_axes，这里就必须加上它。否则TensorRT会认为batch是固定的，动态batch就失效了。

解析成功后，network里就包含了模型的计算图。但这时候还不能直接跑，得先构建Engine。

3.4 Engine序列化与反序列化

Engine是TensorRT优化后的最终产物。构建Engine很耗时，尤其是复杂模型，可能要好几分钟。所以咱们一般会把它序列化保存到磁盘，下次直接加载，省得重新构建。

3.4.1 构建并序列化Engine

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

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

# 保存到文件
with open("resnet18.engine", "wb") as f:
    f.write(serialized_engine)

print("Engine序列化完成，保存为 resnet18.engine")

构建时，set_memory_pool_limit 控制显存使用上限。设得太小，构建可能失败；设得太大，显存不够用。我一般先设1GB，不够再加。

3.4.2 反序列化加载Engine

# 从文件加载
with open("resnet18.engine", "rb") as f:
    serialized_engine = f.read()

# 创建runtime
runtime = trt.Runtime(logger)

# 反序列化
engine = runtime.deserialize_cuda_engine(serialized_engine)

# 创建执行上下文
context = engine.create_execution_context()

print("Engine反序列化成功，准备推理")

反序列化很快，基本是毫秒级的。所以线上服务启动时，直接从文件加载Engine，比重新构建快得多。

3.5 完整流程总结

咱们把整个流程串起来，就是下面这张表：

步骤	操作	输出	耗时
1	PyTorch导出ONNX	.onnx文件	秒级
2	ONNX检查与优化	优化后的.onnx	秒级
3	TensorRT解析ONNX	network对象	秒级
4	构建Engine并序列化	.engine文件	分钟级
5	反序列化加载Engine	engine对象	毫秒级

嗯，到这里，模型导出的基础流程就讲完了。下一章咱们会深入Engine的构建配置，看看怎么调优才能让推理跑得更快。到时候我会分享一些实际项目中的调参经验，保证让你少走弯路。