ONNX与TensorRT模型部署流程

ONNX（Open Neural Network Exchange）和TensorRT是深度模型部署中最核心的两个工具。

ONNX——模型交换格式：

• 由微软和Facebook于2017年联合推出。
• 定义了一套标准的算子集和计算图格式（.onnx文件），使模型可以在不同框架间迁移。
• 支持PyTorch→ONNX→TensorRT/ONNX Runtime/TFLite的转换路径。

PyTorch → ONNX导出流程：

import torch.onnx
torch.onnx.export(
    model,              # PyTorch模型
    dummy_input,        # 示例输入（用于跟踪）
    'model.onnx',       # 输出路径
    export_params=True, # 导出训练好的参数
    opset_version=17,   # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'},
                  'output': {0: 'batch_size'}}  # 动态batch
)

• 注意事项：
  • 使用torch.onnx.export时自动使用Tracing（跟踪执行）。
  • 动态控制流（if/for依赖输入）可能失败，需使用TorchScript。
  • 一些PyTorch算子可能没有ONNX对应映射。

ONNX优化：

• 使用onnx-simplifier（onnxsim）：消除形状推断、常量折叠、冗余节点删除。
• 使用onnxruntime进行推理和精度验证。

TensorRT——推理优化引擎：

• NVIDIA的深度学习推理优化库。
• 支持INT8/FP16/FP32精度。
• 核心优化技术：
  1. 层融合（Layer Fusion）：合并Conv+BN+ReLU等连续操作为单个Kernel。
  2. 权重量化：INT8量化（需要校准数据集）。
  3. 内核自动调优：为每层选择最快的CUDA kernel。
  4. 内存优化：内存池复用。
  5. 动态Tensor：自动处理可变输入维度。

ONNX → TensorRT流程：

import tensorrt as trt

# 方式1：trtexec命令行
trtexec --onnx=model.onnx --saveEngine=model.trt --fp16 --workspace=4096

# 方式2：Python API
builder = trt.Builder(logger)
network = builder.create_network(EXPLICIT_BATCH)
parser = trt.OnnxParser(network, logger)
parser.parse_from_file('model.onnx')

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)
config.set_flag(trt.BuilderFlag.FP16)

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

完整部署流水线： PyTorch模型 → ONNX导出 → ONNX简化(onnxsim) → ONNX→TensorRT构建Engine → TensorRT推理

性能提升：

• FP16：2×加速，几乎无精度损失。
• INT8：4-6×加速（相比FP32），精度损失1-3%。

常见问题：

• Dynamic Shapes在TensorRT中需要设置优化范围（Optimization Profile）。
• 某些算子（如NonMaxSuppression）在TensorRT中实现有限。
• 建议使用TensorRT的Plugin API实现自定义算子。