ONNX与TensorRT推理优化应用实践 - 第1节

1.1 我们为什么需要ONNX与TensorRT

先聊个实际的问题。你辛苦训练好的模型，总不能一直放在实验室里跑吧？到了生产环境，要上服务器、要进手机、要跑到边缘设备上，这时候你会发现——模型推不动了。

为什么会这样？因为训练框架（PyTorch、TensorFlow）本身不是为极致推理性能设计的。它们要的是灵活性，为了反向传播、动态图、调试方便，底层加了很多“中间层”和“检查逻辑”。这些都拖慢了推理速度。

所以我们需要两个东西：

  • 一个通用的模型中间格式——让训练好的模型能跨框架、跨设备迁移。
  • 一个高性能的推理引擎——把模型压榨到硬件极限。

前者是ONNX，后者是TensorRT。我个人习惯把它们看作推理链条上的“翻译官”和“赛车手”。

1.2 ONNX：模型界的“通用语言”

ONNX，全称Open Neural Network Exchange，说白了就是一套标准。它把你的PyTorch模型、TensorFlow模型、甚至Scikit-learn模型，统一成同一个文件格式（.onnx）。

好处很明显：

  • 你不用再为“这个模型只能在这个框架里跑”而头疼。
  • 你可以拿ONNX模型去做跨平台部署——Windows、Linux、ARM、GPU、NPU甚至某些专用芯片都认它。
  • 它内置了一组标准算子（Conv、Relu、MatMul、Softmax等），大多数模型的95%以上操作都能直接映射过去。
我的经验：我在项目中遇到过，有些模型用的算子太冷门了（比如某个自定义的循环结构），ONNX里没有对应算子。这时候有两种解法：一是用ONNX的“自定义算子”机制自己注册，二是把模型拆成几个子图来导出。我个人偏向第二种，灵活度更高。

ONNX的导出流程

ONNX其实不负责做图变换或者算子融合，它只做“格式转换”。一个典型的导出流程是这样的：

# PyTorch -> ONNX 示例
import torch
import torch.onnx

model = torchvision.models.resnet18(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224)

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

注意那个dynamic_axes参数。如果不设置，导出的ONNX模型就只能处理固定batch size的输入。设置之后，线上服务可以灵活改变batch size。嗯，这里要记住。

1.3 TensorRT：把模型装进“硬件冰箱”

TensorRT是NVIDIA专门给自家GPU做推理优化的工具。它能做的事情，不止是“运行模型”，而是把模型“重新编译”成最适合当前GPU架构的执行计划（engine）。

它干了什么？

  • 算子融合：比如把Conv + BN + Relu拼成一个Kernel执行，省掉中间读写。
  • 精度校准：支持FP16、INT8、甚至INT4推理，速度翻倍，精度损失可控。
  • 内存复用：对tensor的生命周期做分析，减少显存碎片和重复分配。
  • 动态shape支持：允许同一个engine处理不同尺寸的输入。
注意：TensorRT不是ONNX的替代品。ONNX是“格式”，TensorRT是“引擎”。你可以直接用ONNX Runtime做推理，也可以把ONNX模型转成TensorRT engine再做推理。后者性能更好，但构建时间也更长（因为要做图优化和Kernel选择）。我在项目里通常把TensorRT model building放到离线阶段，线上只做inference。

TensorRT的典型工作流

# ONNX -> TensorRT Engine 示例
import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("resnet18.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB
config.set_flag(trt.BuilderFlag.FP16)

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

这个serialized_engine就是能在线上直接加载并运行的文件。构建一次，运行多次。

1.4 应用场景：什么时候该用谁？

场景推荐方案原因
跨平台部署（CPU/GPU、Windows/Linux/ARM）ONNX RuntimeONNX RT后端支持多设备，生态最广
NVIDIA GPU上追求极致性能（服务器端）ONNX -> TensorRTTensorRT推理延迟可以比ONNX RT低2~3倍
边缘设备（Jetson系列）TensorRTJetson带TensorRT硬加速单元（DLA、PVA）
模型原型快速验证ONNX Runtime + CPU backend不需要GPU环境也能跑，开发方便
INT8量化部署TensorRT + calibrationTensorRT的INT8量化方案最成熟、加速比最高

你看，没有一个银弹。我经常遇到新人问“哪个更厉害”？其实不是谁更厉害，而是谁更适合你的硬件和你的性能目标。

1.5 软硬件要求与环境搭建

环境配置阶段，我建议你一次配好，后面少踩坑。下面是需要准备的东西：

  • 操作系统：Ubuntu 20.04 / 22.04（Windows也能搞，但坑更多，我建议用Linux）
  • CUDA：11.x 或 12.x（具体版本要和后面工具对齐）
  • cuDNN：与CUDA版本对应，TensorRT 8.x以后大部分依赖内建了
  • TensorRT：推荐8.6+，支持更多ONNX算子
  • ONNX Runtime：1.14+，配合TensorRT Execution Provider使用
  • GPU：至少4GB显存（T4、V100、A10、RTX 3060及以上都行）
版本兼容性要记住：我踩过最大的坑是TensorRT版本和cuDNN版本不匹配。结果build engine时一直报“unsupported memory type”。后来查了NVIDIA的官方兼容矩阵才解决。建议装之前先去NVIDIA官网确认下兼容关系。

安装步骤（Ubuntu + CUDA 11.8为例）

第一步，装CUDA和cuDNN（这里不展开，按NVIDIA官方指南走即可）。

第二步，装TensorRT。可以去NVIDIA官网下载tar包或者用pip:

pip install tensorrt==8.6.1
# 验证安装
python -c "import tensorrt; print(tensorrt.__version__)"

第三步，装ONNX Runtime。建议带上TensorRT Execution Provider:

pip install onnxruntime-gpu==1.15.1
# 或者从源码构建，包含TensorRT支持
pip install onnxruntime-trt==1.15.1

第四步，验证ONNX Runtime是否能找到TensorRT:

import onnxruntime as ort
providers = ort.get_available_providers()
print(providers)
# 预期输出中应包含 'TensorrtExecutionProvider'

为什么我要强调这一步？因为很多人装了TensorRT后，ONNX Runtime默认不知道它。需要设置环境变量或者指定provider列表。我曾经在一个客户现场调试了半天，发现就是没配provider，白白浪费一个下午。

1.6 本章知识脉络

下面这张图概括了我们第一章的核心逻辑：从训练框架到ONNX再到TensorRT，最后部署到GPU上。

ONNX 与 TensorRT 推理优化核心流程 PyTorch / TF 训练好的模型 ONNX 格式 通用中间表示 ONNX Runtime 多后端推理 TensorRT Engine 高性能推理引擎 NVIDIA GPU 线上推理服务 torch.onnx.export trt.OnnxParser 关键决策点 如果追求跨平台兼容性，走 ONNX Runtime 路径；如果追求极致 GPU 性能，转 TensorRT 后部署
再啰嗦一句：不要在构建TensorRT engine时贪图最小的workspace，这会导致引擎在遇到大tensor时直接OOM。我一般给1~2GB workspace，大多数模型够了。如果遇到超大的Batch，可以适当增加。

环境搭好，后面的章节就可以动手实操了。记住：ONNX和TensorRT不是魔法，它们只是帮你把模型调教到硬件该有的状态。学会它们，你的模型才能真正“跑起来”。

← 上一章 📖 返回目录 下一章 →