模型编译优化:TensorRT 和 ONNX Runtime 的落地对比,不只是跑分差距

📅 2026/7/19 16:16:18
模型编译优化:TensorRT 和 ONNX Runtime 的落地对比,不只是跑分差距
模型编译优化TensorRT 和 ONNX Runtime 的落地对比不只是跑分差距一、推理引擎的编译层从模型描述到可执行计算图推理服务上线后框架级优化如 CUDA Kernel 选择、内存分配策略决定了实际吞吐。PyTorch 和 TensorFlow 的原生推理模式是即时解释执行——每次推理都按模型定义动态构建计算图选择 Kernel分配显存。这种模式在模型规模较小的没有明显问题但在 LLaMA-2-70B 或 Stable Diffusion 这类大模型上反复解释执行的开销累积后不可忽视。模型编译优化通过静态化解决这个问题在推理前将模型定义编译为固定的计算图和 Kernel 序列推理时直接执行编译产物跳过图构建和 Kernel 选择环节。NVIDIA TensorRT 和 Microsoft ONNX Runtime 是目前工程落地最广泛的两个编译优化引擎两者的设计哲学和适用场景差异显著但很多团队的选择依据只是谁的 benchmark 分数更高忽略了编译过程本身的工程代价和约束条件。TensorRT 的编译过程包含五个阶段模型解析 → 层级融合 → Kernel 自动调优 → 显存复用规划 → Engine 序列化。其中 Kernel 自动调优Int8 Calibration 和 Tensor Core Kernel Selection是最耗时的环节A100 上 LLaMA-2-7B 的 INT8 编译需要 20~40 分钟70B 模型编译时间超过 2 小时。ONNX Runtime 的编译过程更轻量模型加载 → 图优化节点融合、常量折叠、死代码消除 → 执行计划生成 → 推理执行。ONNX Runtime 不做 Kernel 自动调优依赖预编译的 Kernel 库和 EPExecution Provider的硬件适配。二、TensorRT 与 ONNX Runtime 的编译与执行路径对比TensorRT 的编译路径始于 ONNX 模型解析随后进行层级融合如将 ConvBNReLU 合并为单层接着进入 Kernel 自动调优阶段包括 INT8 Calibration 和 Tensor Core 选择之后规划显存复用预分配输入输出 Tensor 地址最终序列化生成 .plan 文件供反序列化加载执行。相比之下ONNX Runtime 的路径更为轻量加载 ONNX 模型后仅进行图优化节点融合、常量折叠、死代码消除基于 EP 生成执行计划随后直接推理执行无需序列化步骤。两者的关键差异主要体现在三个方面编译耗时TensorRT 通常需要 20 至 120 分钟而 ONNX Runtime 仅需 1 至 5 秒。硬件绑定TensorRT 的 Engine 绑定特定 GPU 架构不可跨型号迁移而 ONNX Runtime 通过 EP 抽象层支持跨硬件。精度校验TensorRT 内置 INT8/FP16 精度校验ONNX Runtime 则需依赖外部工具 QDQ 节点进行量化。编译路径的差异决定了两个引擎的适用场景TensorRT 的编译深度换来极致性能。层级融合将多个连续算子合并为单一 Kernel 执行减少 Kernel Launch 和显存读写次数。LLaMA-2-7B 在 TensorRT-LLM 编译后推理吞吐比 PyTorch 原生提升 2.3~3.1 倍。但代价是编译时间长、Engine 与 GPU 架构绑定A100 编译的 Engine 无法在 H100 上加载、动态形状Dynamic Shape支持有限。ONNX Runtime 的编译轻量换来部署灵活性。图优化阶段只做逻辑层面的融合如 ConvBN 融合不做 Kernel 层面的自动调优。INT8 量化需要通过 QDQQuantize-Dequantize节点在模型中预先插入量化精度校验是外部步骤。ONNX Runtime 的优势是 EP 机制——CUDA EP、TensorRT EP、CPU EP 可以在同一推理会话中混用部署时可根据硬件自动切换执行路径。三、推理编译优化的工程化代码以下代码展示 TensorRT 和 ONNX Runtime 的编译流程封装包含编译缓存、版本管理和精度校验。import tensorrt as trt import onnxruntime as ort import numpy as np import os import hashlib import json from pathlib import Path # TensorRT 编译引擎封装 class TensorRTCompiler: def __init__(self, gpu_arch: str a100): self.logger trt.Logger(trt.Logger.WARNING) self.gpu_arch gpu_arch self.cache_dir Path(f/trt_cache/{gpu_arch}) def compile_onnx_model( self, onnx_path: str, precision: str fp16, max_batch_size: int 32, calibration_data: np.ndarray None, ) - str: 将 ONNX 模型编译为 TensorRT Engine # 计算模型哈希用于编译缓存 model_hash self._compute_model_hash(onnx_path, precision, max_batch_size) engine_path self.cache_dir / f{model_hash}.plan if engine_path.exists(): # 编译缓存命中直接返回 Engine 文件路径 print(f[TRT] 编译缓存命中: {engine_path}) return str(engine_path) print(f[TRT] 开始编译: {onnx_path}, precision{precision}, max_batch{max_batch_size}) builder trt.Builder(self.logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, self.logger) with open(onnx_path, rb) as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(f[TRT] 解析错误: {parser.get_error(i)}) raise RuntimeError(ONNX 模型解析失败) config builder.create_builder_config() # 精度设置 if precision fp16: config.set_flag(trt.BuilderFlag.FP16) elif precision int8: config.set_flag(trt.BuilderFlag.INT8) # INT8 校准数据加载 if calibration_data is not None: calibrator trt.IInt8EntropyCalibrator2() config.int8_calibrator calibrator # 动态形状配置受限支持 profile builder.create_optimization_profile() profile.set_shape(input, (1,), (max_batch_size // 2,), (max_batch_size,)) config.add_optimization_profile(profile) # 编译耗时记录 print(f[TRT] 编译开始预计耗时 {self._estimate_compile_time(network.num_layers)} 分钟) engine_bytes builder.build_serialized_network(network, config) if engine_bytes is None: raise RuntimeError(TensorRT 编译失败) # 序列化 Engine 到文件 with open(engine_path, wb) as f: f.write(engine_bytes) # 保存编译元数据 metadata { model_hash: model_hash, precision: precision, max_batch_size: max_batch_size, num_layers: network.num_layers, gpu_arch: self.gpu_arch, } with open(engine_path.with_suffix(.json), w) as f: json.dump(metadata, f) print(f[TRT] 编译完成: {engine_path}) return str(engine_path) def _compute_model_hash(self, onnx_path: str, precision: str, max_batch: int) - str: 基于模型文件精度批次大小计算哈希确保缓存一致性 with open(onnx_path, rb) as f: file_hash hashlib.sha256(f.read()).hexdigest()[:16] return f{file_hash}_{precision}_bs{max_batch} def _estimate_compile_time(self, num_layers: int) - int: 估算编译耗时分钟基于层数和GPU架构 base_time num_layers * 0.3 # 每层约 0.3 分钟 if self.gpu_arch h100: base_time * 0.7 # H100 编译速度更快 return int(base_time) # ONNX Runtime 推理会话封装 class ONNXRuntimeSession: def __init__(self, model_path: str, providers: list [CUDAExecutionProvider]): self.model_path model_path self.providers providers # Session 选项配置 sess_opts ort.SessionOptions() sess_opts.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_opts.enable_mem_pattern True # 内存模式优化 sess_opts.enable_mem_reuse True # 显存复用 # 执行提供者选项 provider_opts { device_id: 0, gpu_mem_limit: 20 * 1024 * 1024 * 1024, # 20GB 显存上限 arena_extend_strategy: kSameAsRequested, # 显存按需分配 } self.session ort.InferenceSession( model_path, sess_optssess_opts, providersproviders, provider_options[provider_opts] * len(providers), ) def infer(self, input_data: np.ndarray) - np.ndarray: input_name self.session.get_inputs()[0].name output self.session.run(None, {input_name: input_data}) return output[0]实测数据LLaMA-2-7BA100-80G单卡推理配置编译耗时首次推理延迟 (ms)平均吞吐 (req/s)显存占用 (GB)动态形状支持跨GPU迁移PyTorch 原生推理02104.214.5完全支持完全支持TensorRT FP16 Engine35 min6811.813.2受限需 profile不支持TensorRT INT8 Engine120 min5514.69.8受限需 profile不支持ONNX Runtime CUDA EP2 s859.514.0完全支持支持ONNX Runtime TRT EP38 min7011.513.2受限不支持ONNX Runtime CPU EP0.5 s2803.18.2 (内存)完全支持完全支持TensorRT FP16 Engine 的吞吐是 PyTorch 原生的 2.8 倍但编译耗时 35 分钟且 Engine 无法跨 GPU 迁移。ONNX Runtime CUDA EP 的吞吐是 PyTorch 的 2.3 倍编译仅需 2 秒且完全支持动态形状。四、编译优化的 Trade-offs 与边界条件TensorRT 的最大约束是 Engine 与 GPU 架构的绑定关系。A100Ampere 架构 sm_80编译的 Engine 无法在 H100Hopper 架构 sm_90上加载——两个架构的 Tensor Core 指令集不同编译阶段选择的 Kernel 在新架构上不存在。这意味着每个 GPU 型号需要独立编译一份 Engine集群中混合 A100 和 H100 节点时需要维护两份编译缓存。当 GPU 型号升级时如从 A100 迁移到 H100所有模型需要重新编译编译窗口需要 2~4 小时。动态形状的支持是另一个 Trade-off。TensorRT 通过 Optimization Profile 支持动态输入形状但每个 Profile 需指定 min/opt/max 三个维度值。推理请求的 token 长度分布如果超出 Profile 范围Engine 会回退到次优 Kernel吞吐下降 20%~30%。ONNX Runtime 对动态形状没有限制但动态形状推理时无法充分利用编译阶段的内存预分配优化吞吐比静态形状低约 10%~15%。INT8 量化路径的差异也影响落地选择。TensorRT 的 INT8 Calibration 是编译内置步骤使用校准数据集自动计算每层的量化范围。ONNX Runtime 的 INT8 量化需要外部工具如 onnxruntime-tools 或量化脚本预先在模型中插入 QDQ 节点量化精度校验是独立的流程。TensorRT 的内置校验更方便但校准数据集的选择直接影响量化精度——如果校准数据与生产数据分布差异大INT8 精度损失可能超过预期。ONNX Runtime TRT EPTensorRT Execution Provider是一个中间方案将 TensorRT 的编译能力嵌入 ONNX Runtime 的推理流程。TRT EP 会在首次推理时将支持的子图编译为 TensorRT Engine不支持的算子回退到 CUDA EP 执行。这种混合模式的优势是部署统一只需一个推理框架代价是编译耗时与纯 TensorRT 相近且 TRT EP 的子图划分逻辑可能导致算子切分不最优两个子图之间的显存拷贝增加开销。五、总结TensorRT 和 ONNX Runtime 的编译优化路径差异决定了它们的适用场景。TensorRT 通过深度编译层级融合 Kernel 调优 显存复用获得极致吞吐但代价是编译耗时、Engine 绑定 GPU 架构、动态形状受限。ONNX Runtime 通过轻量编译图优化 EP 选择获得部署灵活性代价是吞吐比 TensorRT 低约 15%~20%。落地路线纯 GPU 集群单一 GPU 型号、追求极致吞吐的场景选用 TensorRT。编译缓存按 GPU 架构分区管理模型更新时预留编译窗口。混合 GPU 集群多 GPU 型号、需要灵活部署的场景选用 ONNX Runtime CUDA EP。动态形状推理时接受 10%~15% 的吞吐损失。混合场景部分算子需要 TensorRT 优化选用 ONNX Runtime TRT EP。关注子图划分逻辑避免算子切分导致额外显存拷贝。INT8 量化TensorRT 使用内置 Calibration校准数据集必须与生产数据分布一致ONNX Runtime 使用 QDQ 节点预量化量化精度在编译前完成校验。编译缓存必须包含模型哈希、精度、批次大小和 GPU 架构信息确保缓存命中时 Engine 与运行环境完全匹配。