AI 推理优化实战：ONNX Runtime 与 TensorRT 的性能对比与部署选型

一、模型推理的延迟瓶颈：为什么训练快推理慢

深度学习模型在训练和推理阶段面临截然不同的性能瓶颈。训练阶段关注吞吐量（每秒处理多少样本），可以利用大 Batch Size 和梯度累积来充分利用 GPU；推理阶段关注延迟（单个请求的响应时间），通常 Batch Size 为 1，GPU 利用率极低。一个在训练时吞吐量 1000 samples/s 的模型，推理时单样本延迟可能高达 50ms，远不能满足实时服务的需求。

推理延迟的三大来源：计算量（FLOPs）、内存访问（Memory Bandwidth）和框架开销（Python 解释器、动态图调度）。PyTorch 的动态图机制在训练时提供了灵活性，但推理时每次执行都要重新构建计算图，引入了不必要的开销。此外，PyTorch 的算子实现未针对推理场景做极致优化（如 Kernel 融合、精度混合），导致 GPU 计算单元利用率不足。

二、两种推理优化框架的架构差异

ONNX Runtime 和 TensorRT 是两种主流的推理优化框架，它们的优化策略和适用场景有显著差异。

flowchart TB A[PyTorch 模型] --> B[导出 ONNX] B --> C{推理框架选择} C --> D[ONNX Runtime] C --> E[TensorRT] D --> D1[图优化: 算子融合 + 常量折叠] D --> D2[量化: INT8/FP16] D --> D3[执行提供器: CPU/GPU/TensorRT] D1 & D2 & D3 --> D4[通用性高, 跨平台] E --> E1[图优化: 层融合 + Kernel 自动调优] E --> E2[量化: INT8/FP8 校准] E --> E3[内存优化: 张量复用 + 流水线] E1 & E2 & E3 --> E4[NVIDIA GPU 专属, 极致性能] D4 & E4 --> F[性能基准测试]

ONNX Runtime 是通用推理框架，支持 CPU、GPU 和多种加速器，优化以图级别变换为主；TensorRT 是 NVIDIA 专属推理框架，针对 NVIDIA GPU 做了极致优化，包括 Kernel 自动调优、内存复用和 FP8 支持。

三、两种框架的部署实现与性能对比

3.1 ONNX Runtime 推理部署

""" ONNX Runtime 推理部署 支持 CPU 和 GPU 执行提供器 """ import onnxruntime as ort import numpy as np import torch from transformers import BertModel import time class OnnxRuntimeInferencer: """ONNX Runtime 推理器""" def __init__(self, onnx_path: str, provider: str = "CUDAExecutionProvider"): """ 初始化推理会话 provider: CPUExecutionProvider 或 CUDAExecutionProvider """ # 配置会话选项 sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ( ort.GraphOptimizationLevel.ORT_ENABLE_ALL) sess_options.intra_op_num_threads = 4 # 创建推理会话 providers = [provider, "CPUExecutionProvider"] self.session = ort.InferenceSession( onnx_path, sess_options=sess_options, providers=providers) # 获取输入输出信息 self.input_names = [inp.name for inp in self.session.get_inputs()] self.output_names = [out.name for out in self.session.get_outputs()] def predict(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray: """执行推理""" inputs = { "input_ids": input_ids.astype(np.int64), "attention_mask": attention_mask.astype(np.int64), } outputs = self.session.run(self.output_names, inputs) return outputs[0] def benchmark(self, input_ids: np.ndarray, attention_mask: np.ndarray, num_iterations: int = 100) -> dict: """性能基准测试""" # 预热 for _ in range(10): self.predict(input_ids, attention_mask) # 正式测试 latencies = [] for _ in range(num_iterations): start = time.perf_counter() self.predict(input_ids, attention_mask) latencies.append( (time.perf_counter() - start) * 1000) return { "mean_ms": np.mean(latencies), "p50_ms": np.percentile(latencies, 50), "p99_ms": np.percentile(latencies, 99), "framework": "ONNX Runtime", } def export_to_onnx(model: BertModel, onnx_path: str, seq_len: int = 128): """将 PyTorch 模型导出为 ONNX 格式""" model.eval() dummy_input_ids = torch.randint( 0, 30000, (1, seq_len)) dummy_attention_mask = torch.ones(1, seq_len, dtype=torch.long) torch.onnx.export( model, (dummy_input_ids, dummy_attention_mask), onnx_path, input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={ "input_ids": {0: "batch", 1: "seq_len"}, "attention_mask": {0: "batch", 1: "seq_len"}, "logits": {0: "batch"}, }, opset_version=17, do_constant_folding=True, )

3.2 TensorRT 推理部署

""" TensorRT 推理部署 利用 NVIDIA GPU 的极致优化能力 """ import tensorrt as trt import numpy as np import time class TensorRTInferencer: """TensorRT 推理器""" def __init__(self, engine_path: str): """加载 TensorRT 引擎""" self.logger = trt.Logger(trt.Logger.WARNING) self.runtime = trt.Runtime(self.logger) with open(engine_path, "rb") as f: engine_data = f.read() self.engine = self.runtime.deserialize_cuda_engine( engine_data) self.context = self.engine.create_execution_context() # 分配 GPU 内存 self.inputs = [] self.outputs = [] self.bindings = [] for i in range(self.engine.num_io_tensors): name = self.engine.get_tensor_name(i) shape = self.engine.get_tensor_shape(name) dtype = trt.nptype(self.engine.get_tensor_dtype(name)) size = trt.volume(shape) # 分配 GPU 内存 import pycuda.driver as cuda import pycuda.autoinit mem = cuda.mem_alloc(size * np.dtype(dtype).itemsize) self.bindings.append(int(mem)) if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: self.inputs.append( {"name": name, "mem": mem, "shape": shape, "dtype": dtype}) else: self.outputs.append( {"name": name, "mem": mem, "shape": shape, "dtype": dtype}) def predict(self, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray: """执行推理""" import pycuda.driver as cuda # 将输入数据拷贝到 GPU np.copyto( cuda.mem_alloc_like( self.inputs[0]["mem"]), input_ids.astype(self.inputs[0]["dtype"])) cuda.memcpy_htod( self.inputs[0]["mem"], input_ids.astype(self.inputs[0]["dtype"])) cuda.memcpy_htod( self.inputs[1]["mem"], attention_mask.astype(self.inputs[1]["dtype"])) # 设置输入输出地址 for inp in self.inputs: self.context.set_tensor_address( inp["name"], int(inp["mem"])) for out in self.outputs: self.context.set_tensor_address( out["name"], int(out["mem"])) # 执行推理 self.context.execute_async_v3( stream=cuda.Stream()) # 将输出数据拷贝回 CPU output = np.empty( self.outputs[0]["shape"], dtype=self.outputs[0]["dtype"]) cuda.memcpy_dtoh(output, self.outputs[0]["mem"]) return output def build_tensorrt_engine(onnx_path: str, engine_path: str, fp16: bool = True, int8: bool = False, calibration_data=None): """ 从 ONNX 模型构建 TensorRT 引擎 包含 FP16/INT8 量化优化 """ 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) # 解析 ONNX 模型 with open(onnx_path, "rb") as f: if not parser.parse(f.read()): for error in range(parser.num_errors): print(f"ONNX 解析错误: {parser.get_error(error)}") return # 配置构建器 config = builder.create_builder_config() config.set_memory_pool_limit( trt.MemoryPoolType.WORKSPACE, 4 << 30) # 4GB 工作空间 if fp16: config.set_flag(trt.BuilderFlag.FP16) print("已启用 FP16 量化") if int8 and calibration_data is not None: config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator = calibration_data print("已启用 INT8 量化") # 构建引擎 print("正在构建 TensorRT 引擎（可能需要数分钟）...") engine_bytes = builder.build_serialized_network(network, config) with open(engine_path, "wb") as f: f.write(engine_bytes) print(f"TensorRT 引擎已保存: {engine_path}")

四、两种框架的适用边界与部署陷阱

ONNX Runtime 的算子兼容性：并非所有 PyTorch 算子都能导出为 ONNX 格式。自定义算子、动态控制流（如if/else依赖输入值）和部分高级操作（如torch.nn.functional.grid_sample的某些模式）可能导致导出失败或精度偏差。建议在导出后用onnxruntime的验证模式对比 PyTorch 和 ONNX 的输出差异，确保精度损失在 1e-5 以内。

TensorRT 的构建耗时：TensorRT 引擎构建时会对每层做 Kernel 自动调优，在当前 GPU 上测试不同 Kernel 实现的性能，选择最优方案。这个过程可能需要 10-30 分钟，且构建结果与 GPU 型号绑定——在 A100 上构建的引擎无法在 V100 上使用。生产环境中建议在部署目标 GPU 上预构建引擎，将构建产物（.engine 文件）纳入版本管理。

动态形状的支持差异：ONNX Runtime 对动态形状（如可变序列长度）的支持较好，通过dynamic_axes配置即可。TensorRT 对动态形状的支持通过 Optimization Profile 实现，需要指定每个维度的最小值、最优值和最大值，且运行时形状变化会触发重新调优，增加延迟。对于序列长度变化频繁的场景，建议使用固定长度 + Padding 策略。

INT8 量化的校准数据：TensorRT 的 INT8 量化需要校准数据集来确定量化参数。校准数据应代表生产环境的真实数据分布，如果校准数据与实际数据分布差异大，量化后精度损失可能超过 5%。建议使用 500-1000 条生产数据作为校准集，并在量化后用完整验证集评估精度。

五、总结

推理优化框架的选型核心在于"通用性 vs. 极致性能"的权衡。ONNX Runtime 通用性强、跨平台、部署简单，适合快速上线和多平台场景；TensorRT 在 NVIDIA GPU 上性能极致，适合延迟敏感的在线服务。落地时建议先用 ONNX Runtime 快速验证，确认精度和延迟基本满足需求后，再考虑迁移到 TensorRT 做极致优化。FP16 量化是成本最低的优化手段（精度损失通常 < 0.5%），应优先启用；INT8 量化需要校准数据，建议在 FP16 不满足需求时再考虑。