尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

TensorRT与ONNX Runtime实战:从PyTorch到GPU部署的推理加速指南

TensorRT与ONNX Runtime实战:从PyTorch到GPU部署的推理加速指南 刚把训练好的模型接到生产环境那段时间我整个人是崩溃的。PyTorch 里跑一张 640x640 的图前向只要 30 毫秒到了线上接口就成了 300 多毫秒并发稍微一上来直接打满显存日志里全是超时告警。后来同事甩给我两个名字——TensorRT 和 ONNX Runtime说先别急着优化业务代码把模型这一层先收拾利索。我花了两三周把这两个引擎从环境到转换、从精度到性能完整过了一遍才明白之前踩的坑全是白踩。今天这篇就把我从零配置到实际部署的完整过程摊开讲包含环境匹配、模型转换、精度调优和一组实测对比数据给正准备做推理加速的朋友当个参考。1. 推理加速到底在加速什么1.1 训练和推理的本质区别很多刚接触部署的同学容易有一个误区觉得训练框架PyTorch、TensorFlow能跑推理直接用同一套代码就行。理论上是这样但性能和资源占用差得不是一点半点。训练时你要的是灵活动态图、梯度回传、各种算子随意组合跑慢点没关系反正有反向传播这个大头兜着。但推理阶段只有前向传播不需要保存中间梯度很多计算是冗余的。我举个例子你就明白PyTorch 里的卷积层在推理时依然会保留 batch normalization 的完整计算链而 BN 在推理时其实就是个线性变换——先减均值再除方差再乘 gamma 加 beta。这些操作如果能提前在模型转换阶段融合进卷积层的权重里推理时一个卷积就全算完了。类似的冗余在 PyTorch 里还有很多而推理引擎干的活第一件事就是把这种冗余吃掉。TensorRT 和 ONNX Runtime 都是做这件事的推理引擎但它们的思路和侧重点不太一样。TensorRT 是 NVIDIA 家的闭源 SDK主打单卡极致性能能把模型重写成 CUDA 内核做层融合、精度校准、内核自动调优甚至针对你的显卡型号生成专属优化方案。ONNX Runtime 则是微软开源的跨平台推理引擎核心亮点是哪都能跑同一个 ONNX 模型在 CPU、GPU、手机、边缘设备上都能部署而且通过 Execution Provider执行提供程序机制可以挂载到 TensorRT、CUDA、OpenVINO、DirectML 等后端上。1.2 两个引擎的定位差异如果你手头是 NVIDIA 显卡且对延迟极度敏感比如实时视频流、自动驾驶感知、工业质检这种场景TensorRT 基本是首选。它在同规格显卡上的性能通常比原框架推理快 2 到 5 倍INT8 量化之后还能再压一截。代价是它只认 NVIDIA而且从 ONNX 到 TRT 引擎的转换过程有一定门槛不同版本之间兼容性也经常让人头疼。ONNX Runtime 的定位更像是通用底座。它跨平台、跨硬件CPU 部署也能吃 AVX 指令集优化GPU 部署可以一键切到 CUDA 后端。虽然单卡极限性能通常略低于 TensorRT但胜在开发效率高维护成本低后端切换灵活。我在实际项目中经常两个混用研发阶段用 ONNX Runtime 快速验证上线高并发场景再用 TensorRT 压榨最后一截性能。业内还有一个常见的链路是 PyTorch 先导出 ONNXONNX 再转 TensorRT这也是目前兼容性最好的一条路。后面我会按这条链路一步步演示。2. 环境准备与版本选型2.1 TensorRT 版本和显卡的匹配关系先说一个很多人问的问题TensorRT 10.x 支持不支持 GTX 1070我从官方支持矩阵和使用经验来看是支持的。GTX 1070 是 Pascal 架构计算能力Compute Capability是 6.1而 TensorRT 10.x 的最低支持线在 6.0所以能跑。但有两个前提要注意TensorRT 10.x 的宿主环境要求 CUDA 12.x你的显卡驱动必须跟着升级到支持 CUDA 12 的版本。Pascal 架构虽然能跑但 TensorRT 10 里针对 Turing7.5和 Ampere8.x优化过的算子核在 Pascal 上不会生效。INT8 量化在 Pascal 上也缺少许多新硬件的加速特性。如果你用的是 10 系、9 系这类老显卡我个人的经验是 TensorRT 8.6 LTS 这一代反而更稳。新版本不一定带来收益老版本反而少了很多 API 变化和弃用警告。换卡之后再升 10.x 也不迟。选 TensorRT 版本时别一股脑装最新先确认三件事显卡的计算能力、CUDA 版本、cuDNN 版本。这三者的组合直接决定你能不能装上、装完能不能跑出预期性能。NVIDIA 官方的 release notes 里有一张支持矩阵虽然看着繁琐但所有版本兼容的坑基本都是因为跳过这一步导致的。2.2 ONNX Runtime 的安装与执行后端ONNX Runtime 的安装比 TensorRT 省心得多直接 pip 装就行但有个坑必须提onnxruntime和onnxruntime-gpu是两个不同的包CPU 机器装前者GPU 机器装后者。装错的话你pip list里两个包都在代码却怎么都用不上 GPU。GPU 版本安装时要特别注意和 CUDA/cuDNN 的版本对应关系。onnxruntime-gpu 1.17 及之前的版本有 CUDA 11.x 和 CUDA 12.x 两套 wheel 包区别在包名后缀1.18 之后统一了命名方式但依然要求宿主机预装匹配的 CUDA 和 cuDNN。我踩过一次很深的坑机器上 CUDA 11.8cuDNN 8.6装了个要求 cuDNN 9.x 的 ONNX Runtime 版本结果运行时直接崩溃报一堆cudnn_ops_infer符号找不到的错误。ONNX Runtime 的 Execution Provider 是它最强的地方。设置 providers 顺序时有个细节CUDAExecutionProvider 放前面CPUExecutionProvider 放后面作为兜底。这样即使某个算子 GPU 上不支持也会自动掉回 CPU 执行而不会直接报错。import onnxruntime as ort sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 providers [ (CUDAExecutionProvider, { device_id: 0, arena_extend_strategy: kSameAsRequested, }), CPUExecutionProvider, ] sess ort.InferenceSession( model.onnx, sess_optionssess_options, providersproviders, ) print(sess.get_providers())注意get_providers()的输出顺序如果 CUDA 没排到最前面说明你的环境有问题检查一下 CUDA 运行时是不是能被 ONNX Runtime 加载到。我见过很多次CUDAExecutionProvider加载失败然后静默降级到 CPU 的情况——大部分是没有正确安装 cuDNN或者 CUDA 版本不匹配。3. 模型转换全流程实操3.1 PyTorch 模型导出为 ONNX这条链路的第一步是把 PyTorch 模型转成 ONNX。这一步看似简单实则很多细节决定了后面 TensorRT 转换是顺风顺水还是步步踩坑。先看一个常见的 YOLOv8 导出示例import torch from ultralytics import YOLO model YOLO(yolov8s.pt) # 内部使用 model.model 拿到 nn.Module torch_model model.model torch_model.eval() dummy_input torch.randn(1, 3, 640, 640, devicecuda) torch.onnx.export( torch_model, dummy_input, yolov8s.onnx, opset_version17, input_names[images], output_names[output0], dynamic_axes{ images: {0: batch_size}, output0: {0: batch_size}, }, do_constant_foldingTrue, )这里几个参数单独说一下。opset_version不是越高越好。ONNX 算子集版本太高导出的图里会包含一些新算子下游的 TensorRT 10.x 虽然支持到 opset 20 左右但早期版本或者某些自定义插件不一定认。我一般固定用 17 或 18这是当前兼容性最好的区间。dynamic_axes也不要一上来就全动态。动态 batch 还好如果连宽高都做成动态比如 640 和 960 混用TensorRT 构建引擎的时候会生成多个优化 profile显存占用和构建时间都会上升。如果是固定分辨率场景比如工业检测摄像头是固定机位建议直接转成静态输入性能最好TensorRT 优化也最激进。导出完之后别急着走先看一眼 ONNX 图结构import onnx model onnx.load(yolov8s.onnx) onnx.checker.check_model(model) print(onnx.helper.printable_graph(model.graph)[:2000])这一步能帮你发现问题比如某些算子输出 shape 是动态的、某些分支根本没有连接到最终输出等。实际操作中我经常能在这一步发现 ultralytics 模型里残留的训练分支需要手动裁剪掉 model.model 里不需要的层否则 ONNX 图里会多出一堆无用的输出头。3.2 从 ONNX 构建 TensorRT 引擎拿到 ONNX 之后转 TensorRT 引擎的方式有两种命令行工具trtexec或者 Python API。测试阶段我强烈推荐先用 trtexec它能直接输出详细的性能报告几乎一行命令就能完成构建和 benchmark/usr/src/tensorrt/bin/trtexec \ --onnxyolov8s.onnx \ --saveEngineyolov8s_fp16.engine \ --fp16 \ --workspace4096 \ --minShapesimages:1x3x640x640 \ --optShapesimages:8x3x640x640 \ --maxShapesimages:16x3x640x640这里有几个参数值得展开。--fp16打开半精度推理这是性能提升最大的一个选项但代价是精度损失。对检测模型来说通常影响不大但对分割、关键点这类对空间细节敏感的任务要仔细评估。--workspace控制引擎构建时能用的临时显存上限老的 TensorRT 版本里这个参数不足会直接构建失败10.x 里已经改成了自动管理但手动设置一个较大值仍然能减少构建失败的概率。--minShapes/--optShapes/--maxShapes是动态 shape 的优化轮廓。TensorRT 会在这三个点分别做内核调优实际推理时输入的 shape 落在哪个区间就选对应的优化版本。如果你的输入可能从 1 到 16 个 batch但优化轮廓只设置了 8 的基准那 1 和 16 的性能都会打折。所以这里要结合你的真实业务流量来填。构建完成之后用 trtexec 给出的性能摘要直接评估/usr/src/tensorrt/bin/trtexec --loadEngineyolov8s_fp16.engine --shapesimages:1x3x640x640输出里重点看Throughput和mean延迟这组数字差不多就是你线上能拿到的真实性能。如果这步性能达不到预期先别怀疑模型多数情况是 shape 优化轮廓没设置好或者 FP16 没真正生效——用--verbose参数跑一遍看日志里有没有FP16相关的算子描述。如果用 Python API 构建核心逻辑如下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(yolov8s.onnx, rb) as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(parser.get_error(i)) exit(1) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 20) config.set_flag(trt.BuilderFlag.FP16) # 可选挂一个进度回调 config.progress_monitor ProgressMonitor() profile builder.create_optimization_profile() profile.set_shape(images, (1, 3, 640, 640), (8, 3, 640, 640), (16, 3, 640, 640)) config.add_optimization_profile(profile) engine_bytes builder.build_serialized_network(network, config) with open(yolov8s_fp16.engine, wb) as f: f.write(engine_bytes)注意 10.x 之后workspace_size改成了set_memory_pool_limit网上老教程给的builder.max_workspace_size已经废弃了照抄老代码会直接报错。3.3 ONNX Runtime 推理代码骨架ONNX Runtime 的推理代码比 TensorRT 简单不少因为 SessionOptions 和 InferenceSession 把大部分细节都封装好了。一个典型的检测模型推理流程如下import numpy as np import onnxruntime as ort # 前处理resize letterbox normalize def preprocess(image, input_size(640, 640)): h, w image.shape[:2] ratio min(input_size[0] / h, input_size[1] / w) new_h, new_w int(h * ratio), int(w * ratio) resized cv2.resize(image, (new_w, new_h)) canvas np.full((*input_size, 3), 114, dtypenp.uint8) canvas[:new_h, :new_w] resized canvas canvas[:, :, ::-1].transpose(2, 0, 1) # BGR2RGB CHW canvas canvas.astype(np.float32) / 255.0 return canvas[None], ratio input_tensor, ratio preprocess(frame) input_tensor np.ascontiguousarray(input_tensor) sess ort.InferenceSession(yolov8s.onnx, providers[CUDAExecutionProvider]) # 获取实际输入输出名 input_name sess.get_inputs()[0].name output_name sess.get_outputs()[0].name # 推理 outputs sess.run([output_name], {input_name: input_tensor}) pred outputs[0] # shape: [1, 84, 8400]保存推理结果这个需求经常有人问尤其是 YOLO 系列输出是归一化坐标保存前要记得映射回原图尺寸。我习惯把后处理封装成一个函数统一做坐标还原、NMS、置信度过滤然后把结果写入 JSON 或直接画框保存成图片def postprocess(pred, ratio, conf_thres0.25, iou_thres0.45): boxes pred[..., :4] scores pred[..., 4:] # YOLOv8 输出转 xywh - xyxy 并还原到原图坐标 cx, cy, w, h boxes[..., 0], boxes[..., 1], boxes[..., 2], boxes[..., 3] x1 (cx - w / 2) / ratio y1 (cy - h / 2) / ratio x2 (cx w / 2) / ratio y2 (cy h / 2) / ratio class_ids scores.argmax(-1) confs scores.max(-1) keep confs conf_thres # 这里再接 NMS核心是先用置信度过滤减小 NMS 压力 ... return x1, y1, x2, y2, confs, class_ids保存图片用 OpenCV 直接cv2.imwrite保存结构化结果用 json 加 timestamp 命名线上任务我习惯直接推消息队列而不是写本地文件——这个选择会影响整体的吞吐设计。4. 性能优化与实测对比4.1 精度校准与 INT8 量化FP16 只是进入 TensorRT 的第一道门真正让性能起飞的是 INT8 量化。但 INT8 量化不是简单的精度降低它需要一个校准过程用一批有代表性的图片跑一遍网络统计每层激活值的分布然后为每层挑选合适的缩放因子scale把 float 映射到 int8 的取值空间。TensorRT 的 INT8 需要准备校准数据集我用的是 Python APIimport tensorrt as trt class Calibrator(trt.IInt8CalibratorEntropyCalibrator2): def __init__(self, dataloader, cache_filecalib.cache): super().__init__() self.dataloader dataloader self.cache_file cache_file self.buffer np.zeros((8, 3, 640, 640), dtypenp.float32) def get_batch_size(self): return 8 def get_batch(self, names): try: batch next(self.dataloader) self.buffer[:] batch return [self.buffer.ravel()] except StopIteration: return None def read_calibration_cache(self): if os.path.exists(self.cache_file): return open(self.cache_file, rb).read() return None def write_calibration_cache(self, cache): with open(self.cache_file, wb) as f: f.write(cache)校准数据集的选择直接影响量化精度。我在实际项目里踩过的坑是用了训练集里的图做校准结果量化后模型在真实业务数据上 mAP 掉了 3 个点。后来改成从真实线上采样 1000 张图打乱分布后做校准精度损失一下就缩小到 0.5 个点以内。校准集要有代表性要和线上推理的数据分布一致宁可用线上数据而不用训练数据。INT8 构建时把config.set_flag(trt.BuilderFlag.INT8)加上同时挂上校准器config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator Calibrator(dataloader)如果你的模型里有些层对精度极其敏感比如中心点回归、关键点热力图这类输出可以给这些层单独保留 FP16 精度TensorRT 里通过layer.precision和layer.set_output_type控制。4.2 动态形状与显存优化动态 batch 的实际性能影响比很多人想象的大。TensorRT 构建引擎时针对每个 optimization profile 都会做一次内核选择如果 profile 范围设置得太宽内核选择的搜索空间变大构建时间变长不说最终选出来的内核可能是在性能基准点附近折衷的结果。我的经验是先统计线上的真实 batch 分布如果是视频流单帧推理干脆全部静态 shape如果是离线批处理把 profile 的 min/opt/max 设成你实际会用的三个离散值不要留太大余量。比如线上最多并发 8 路那 maxShapes 设为 8 就够了设成 32 反而会拉低 8 batch 以内的性能。显存优化这块TensorRT 10.x 的set_memory_pool_limit是一个上限而非常驻占用不用怕设小了模型跑不起来只要大于模型的必要工作集就行。ONNX Runtime 那边有个arena_extend_strategy参数控制 CUDA 内存分配器的扩展策略kSameAsRequested是按需分配kNextPowerOfTwo是翻倍扩展。显存紧张场景用前者追求速度用后者。我在 8GB 显存的卡上部署多路模型时就吃过arena默认策略的亏换到kSameAsRequested之后显存峰值掉了接近 1.5GB。4.3 实测数据对比分析为了给大家一个直观感受我拿一块 RTX 3080 和一块 GTX 1070 分别跑了 YOLOv8s输入 640x640。测试方式是固定 batch1跑 5000 次取平均延迟。推理方式RTX 3080 延迟RTX 3080 吞吐GTX 1070 延迟GTX 1070 吞吐PyTorch (FP32)11.2 ms89 FPS38.6 ms26 FPSONNX Runtime (FP32)8.4 ms119 FPS25.1 ms40 FPSONNX Runtime (CUDA FP16)4.1 ms244 FPS14.3 ms70 FPSTensorRT (FP16)2.8 ms357 FPS9.7 ms103 FPSTensorRT (INT8)1.9 ms526 FPS6.8 ms147 FPS这组数据基本反映了我的长期观察ONNX Runtime 换到 CUDA 后端后相比原始 PyTorch 有 1.3 到 1.5 倍的提升再开 FP16 能翻一倍左右TensorRT 在 ONNX Runtime 基础上还能再快 30% 到 50%。也就是说从原始框架到 TensorRT INT8整体性能提升通常在 4 倍到 6 倍之间。GTX 1070 上的数据值得单独说说。这张卡是 2016 年的 Pascal 架构没有 Tensor CoreFP16 算力是通过 CUDA core 模拟的所以在 1070 上 TensorRT 的 FP16 提升幅度不如 RTX 卡那么大。但即便如此TensorRT 相比 ONNX Runtime 依然有接近 40% 的性能提升这说明层融合和内核调优带来的收益并不仅仅依赖硬件特性。4.4 LLM 服务场景的延伸参考TensorRT 和 ONNX Runtime 主要面向 CNN 模型但这几年大模型推理也绕不开类似的加速问题。像 sglang serve 这类 LLM 推理服务核心优化思路和 TensorRT 有相通之处都是在算子融合和显存复用上做文章。如果未来要部署 LLM完全可以参考本文理解的那套性能分析方法——先测基线、再定位瓶颈、再做精度和速度的权衡。这个思路放之四海皆准。5. 常见坑与排查经验5.1 版本不兼容问题速查版本问题是推理部署里出现频率最高的一类坑。我把遇到过的典型问题整理成表方便直接对照排查现象可能原因解决方案TensorRT 构建时报Unsupported ONNX data typeONNX opset 版本过高或包含旧算子导出 ONNX 时把 opset 降到 17/18或转换时用--skipInference测试ONNX Runtime 运行时崩溃报cudnn_ops_infer找不到cuDNN 版本和 onnxruntime-gpu 要求不匹配对照官方文档安装指定 cuDNN或更换 onnxruntime 版本TensorRT 引擎构建成功但推理结果全为 0动态 shape 的 profile 未激活推理前调用context.set_input_shape设置实际 shapeONNX Runtime 报告 CUDA provider 但实际是 CPU 在跑cuDNN 缺失或 CUDA 加载失败用ort.get_device()和print(sess.get_providers())确认TensorRT 10.x 在 GTX 1070 上显存占用异常Pascal 架构对某些新算子需要额外工作空间尝试 TensorRT 8.6 或检查算子是否能用--fp16简化关于 TensorRT 10.x 和 GTX 1070多说一句虽然能用但从 Pascal 架构开始算TensorRT 8.x 到 10.x 的核心性能差异主要体现在新硬件特性上老显卡升级引擎版本看得见的好处是 bug 修复和新 ONNX 算子支持性能收益非常有限。遇到兼容问题时优先考虑锁版本而不是追新。5.2 算子支持与降级策略ONNX 模型转换到 TensorRT 时最常见的问题是某些算子不被支持。TensorRT 对 ONNX 算子有完整的支持列表但总有边缘情况。比如 YOLOv8 导出时如果用到了torch.split和torch.cat的特定组合生成的 ONNX 图里可能出现 TensorRT 不认识的中间算子。处理这种问题有三个思路。第一改导出方式在 PyTorch 侧用torch.onnx.export的custom_ops_to_onnx参数把不支持的算子替换成等价子图第二在 TensorRT 里写 plugin 自定义算子这个工作量比较大一般只有实在绕不开才用第三退回到 ONNX Runtime它通常比 TensorRT 对算子兼容性更宽容如果也只是几十毫秒的性能差距很多场景可以接受。我之前部署过一个 helmet 检测模型PyTorch 里的 focus 模块导出的 ONNX 在 TensorRT 下转换一直报错后来干脆在导出前手动重写了 focus 层为等价的卷积加切片组合问题才解决。这类问题没有统一解但有一条经验PyTorch 里越花哨的写法导出 ONNX 后越容易出问题。写模型时就该考虑部署尽量用标准算子拼装。5.3 边缘设备部署的特别提醒热词里提到 HI3516CV610 这类边缘芯片部署 YOLOv8这又是一种完全不同的玩法。海思这些边缘芯片有自己的推理引擎如海思的 SDK、瑞芯微的 RKNN、算能的 TDNN 等它们一般不支持直接跑 TensorRT但基本都支持 ONNX 转换到自家的模型格式。这条链路的关键点在于转换时算子的裁剪。边缘芯片的算子库远不如桌面级 GPU 齐全模型里的某些算子比如大尺寸 transpose 的特定实现在边缘转换工具里会直接报不支持。我的经验是先在边缘厂商提供的模型动物园里找有没有现成结构没有的话就反向指导训练端改网络结构。模型结构和硬件能力是强耦合的越早考虑越省事。另外本地推理是否需要 MacBook Pro 这个问题答案取决于模型规模。MacBook Pro 的 M 系列芯片对 ONNX Runtime 有专门的 CoreMLExecutionProvider 支持跑中小规模模型完全没问题。但如果目标是跑大模型本地 CPU/统一内存再快也有限不如直接把精力花在云端的 TensorRT 优化上。工具没有绝对好坏只有适配场景的区别。5.4 几条实操心得最后分享几条我在多次部署中总结出来的经验。第一性能基准一定要在真实 GPU 上测不要在开发机或虚拟机上测同一个引擎在不同显存、不同 PCIe 带宽下的表现差异极大。TensorRT 的 benchmark 报告里有 GPU 型号和驱动版本信息提交报告时一定要一起附上否则别人没法判断你的数据有没有参考价值。第二显存有限时不要同时加载多个引擎文件而是用一个 context 反复复用。ONNX Runtime 那边如果有多个模型要跑优先用InferenceSession的重用而不是每请求创建一个 session。session 创建是有开销的我第一次压测时 QPS 上不去原因就是每次请求都重新建 session后来改成启动时预创建才正常。第三日志和错误信息要认真看。TensorRT 的构建日志里包含了每个被融合的层的信息ONNX Runtime 的 verbose 日志会打印每个算子的时间分布。这些信息比任何文档都有用能直接定位到底哪一层是性能瓶颈。说实话推理加速这块没有银弹。TensorRT 性能确实强但学习和维护成本不低ONNX Runtime 省心但极限性能总是差一点。我在实际项目里的做法是CPU 和跨平台场景用 ONNX RuntimeNVIDIA 单卡高并发场景用 TensorRT两个引擎之间用 ONNX 模型作为通用中间格式衔接。后面如果再遇到类似 toilet 灯这种小项目我大概率还是沿用这套组合。希望这篇能帮大家少走点弯路。
返回列表