☰
TensorRT原生引擎实战:从ONNX Runtime到多路视频推理的加速优化
2026/10/1 17:36:18 网站建设 项目流程

开头

说起推理加速这件事,我踩过的坑可能比大部分人都要多。最开始我从 PyTorch 直接转 ONNX Runtime,觉得速度已经够用了,后来在一次实时视频流项目中,单路 YOLO 检测在 T4 上跑 720p 输入勉强到 60fps,但对接 8 路视频流时 CPU 占用飙到 90% 以上,GPU 利用率却只有 40% 左右,那一刻我才意识到 ONNX Runtime 的默认优化根本不够。于是我开始研究 TensorRT 原生引擎,从模型转换、精度校准、动态形状到多路并发调度,花了差不多三周时间把整个流程跑通,最后在同样的 T4 上实现了 8 路 1080p@25fps 的实时检测,GPU 利用率稳定在 85% 左右,CPU 占用降到 30% 以下。这篇文章就是我整个探索过程的完整记录,从 ONNX Runtime 的基础加速原理讲起,到 TensorRT 原生的层融合、量化校准和动态 shape 处理,每一步都附上我实际用过的代码、参数和踩过的坑,希望能帮那些正在从"能跑"走向"跑得快"的同行少走点弯路。无论你是刚入门推理部署的新手,还是已经在生产环境里被延迟问题折磨的老手,这篇内容应该都能给你一些参考。

1. 为什么从 ONNX Runtime 转向 TensorRT

1.1 两种推理引擎的定位差异

很多朋友对 ONNX Runtime 和 TensorRT 的理解停留在"都是一个推理框架"这个层面,其实两者干的活完全不同。ONNX Runtime 更像是一个"通用翻译器",它的目标是把各种框架导出的 ONNX 模型(不管是 PyTorch 转的、TensorFlow 转的还是 PaddlePaddle 转的)都能快速跑起来,并且提供一套相对通用的优化手段,比如算子融合、图优化、内存复用等。这些优化是"平台无关"的,它不会针对某一块具体的 GPU 做特别激进的定制。

TensorRT 则完全不一样,它是英伟达把 GPU 硬件架构吃透之后做出来的"专属加速器"。它会针对每一代 GPU 的算子特性、显存带宽、并行调度方式做极其细粒度的优化。比如 Volta 之后的 GPU 都有 Tensor Core,TensorRT 能自动把卷积、全连接这类计算密集的算子转换成 Tensor Core 能执行的指令序列,在保持功能不变的前提下把吞吐量拉高一个量级。ONNX Runtime 在 GPU 上也有 TensorRT 执行提供程序,但那个只是把模型托管给 TensorRT 去做推理,中间多了一层封装和格式转换,性能和直接用 TensorRT 原生引擎还是有差距。

我自己的实测对比是这样的:用 YOLOv5s、640x640 输入、FP16 精度,在 T4 上跑 ONNX Runtime CUDA 执行提供程序,单路延迟稳定在 12ms 左右;同一台机器换成 TensorRT FP16 引擎后,单路延迟能压到 7ms 以内。这还只是单路场景,多路并发时差距更明显,TensorRT 可以通过批量推理(多张图合并成一个 batch 塞进去)获得额外的吞吐量提升,而 ONNX Runtime 的动态 batch 支持就没这么顺手。

所以我的建议是:如果你的项目是“能跑就行,追求兼容性和快速上线”,ONNX Runtime 足够;但如果你做的是视频分析、自动驾驶感知、边缘智能盒子这类对延迟和吞吐量有硬性要求的场景,TensorRT 原生引擎几乎是必经之路。

1.2 TensorRT 加速的核心原理:层融合与精度校准

TensorRT 为什么能有这么大的性能提升,主要是三张王牌:层融合、精度校准、内核自动调优。

先说层融合。GPU 上每个算子执行时都要经历"启动内核->读写数据->输出结果"这个过程,数据在显存里来回搬是很贵的操作。TensorRT 会把多个连续的算子合并成一个,比如卷积后面的批量归一化层和 ReLU 激活层,在传统推理引擎里是三个独立的操作,TensorRT 会把它融合成一个 CBR(Convolution+Bias+ReLU)操作,数据在寄存器里就完成流转了,不用反复读写全局显存。这个融合逻辑在不同 GPU 平台上表现不一样,TensorRT 在编译引擎时就已经根据目标 GPU 的架构做了针对性融合决策,这是 ONNX Runtime 做不到的。

再说精度校准。TensorRT 支持 FP16 和 INT8 两种低精度推理。FP16 对精度的影响通常很小,因为只是在计算过程中把浮点数从 32 位缩到 16 位,动态范围还是够用的。INT8 则需要做校准(Calibration),它会用一批有代表性的数据去统计每层激活值的分布,然后计算一个合适的缩放因子,把浮点数映射到整数范围,这一步做得好能实现接近 3-4 倍的性能提升,做不好会出现明显的精度掉点。我第一次做 INT8 的时候就因为校准数据选得不均匀,导致小目标检测率直接掉了 20%,后来换成包含各种场景光线的测试集才恢复正常。

最后是内核自动调优。TensorRT 在编译引擎时会针对每个算子做大量的 kernel 性能测试,同一层可能生成 10-20 种不同实现,逐个跑一遍选出最优的。这个调优过程很耗时,一份模型在 Jetson 设备上编译可能要 15-20 分钟,但换来的就是在该硬件上的极致性能。这也是为什么 TensorRT 的引擎文件是"设备绑定"的,你在一台机器上编译生成的 .engine 文件换到另一台不同显卡的机器上可能直接加载失败。

2. 环境准备与版本选型

2.1 CUDA、TensorRT、PyTorch 的版本对应关系

这一节我踩过的坑是最多的。TensorRT 对 CUDA 和显卡驱动的版本要求非常严格,不是说你装个最新版就完事了,装完跑起来各种报错,折腾几天才发现是版本不匹配。这里我先给你一个我实测可用的组合,目前在生产环境跑得非常稳定:

组件版本备注
显卡驱动470.63.01建议 >=460.x,低于这个会出现兼容性问题
CUDA11.4不需要装全套,TensorRT 运行只需要 CUDA runtime 的动态链接库
cuDNN8.2.2和 CUDA 11.4 配套
TensorRT8.2.38.2.x 系列对 ONNX 的支持已经比较成熟
Python3.83.6/3.7 也行,但 3.9 以上部分依赖包容易出问题
PyTorch1.10.0导出 ONNX 用,版本别太新,1.10 及以上都行

这个组合是我参考了自己实际环境和主流部署文档整理出来的,比较保守,你不一定要完全照搬,但基本逻辑是把 TensorRT 的版本压到和 CUDA 大版本一致,不要盲目追新。TensorRT 8.4 之后虽然对动态 shape 支持更好了,但我实测下来 8.2.3 在稳定性上反而更靠谱,8.5 有过一个 batch 维度推导的坑,当时查了一整天,最后降级解决。

安装 TensorRT 的时候有个容易忽略的细节:它需要一个专门账号去英伟达官网下载安装包,你如果在公网环境里不方便,可以换一个思路,用 pip 安装 nvidia-pyindex 源里的 tensorrt 包。我这边实测下来 pip 安装的 TensorRT 8.2.3 和官网下载的 tar 包功能完全一致,只是少了配套的 trtexec 工具和 Python 示例代码。trtexec 后面聊转换的时候会讲到,非常有用,建议还是想办法把官方包整个下载下来。

2.2 环境变量与链接库配置

如果你用的是 tar 包方式安装 TensorRT,解压之后需要做的三件事:把 lib 路径加入 LD_LIBRARY_PATH、把 Python 工具包安装进当前 Python 环境、验证导入是否成功。

# 我习惯把 TensorRT 解压到 /opt 下 tar -xzvf TensorRT-8.2.3.0.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz mv TensorRT-8.2.3.0 /opt/ # 把库路径写进 ~/.bashrc,注意别重复追加 echo 'export LD_LIBRARY_PATH=/opt/TensorRT-8.2.3.0/lib:$LD_LIBRARY_PATH' >> ~/.bashrc echo 'export PATH=/opt/TensorRT-8.2.3.0/bin:$PATH' >> ~/.bashrc source ~/.bashrc # 安装 Python API cd /opt/TensorRT-8.2.3.0/python pip install tensorrt-8.2.3.0-cp38-none-linux_x86_64.whl # 验证安装 python -c "import tensorrt as trt; print(trt.__version__)"

这里有个很有意思的细节:TensorRT 的 Python API 和 C++ API 底层是同一套 C++ 库,所以不管你用哪个语言编写推理代码,最终执行效率是一模一样的。Python 版只是给你提供了一套更友好的封装,方便你快速调通流程。我之前遇到过有人担心 Python 调用 TensorRT 会损失性能,这个担心完全可以放下。

验证导入的时候如果报 cuda 相关动态库找不到,大概率是 CUDA 的 lib64 路径没进 LD_LIBRARY_PATH,加上去再验证一次就行。我遇到过更隐蔽的问题:系统里装了多个版本的 CUDA,TensorRT 加载的时候选错版本,直接报 "libcudnn.so.8: cannot open shared object file"。排查方法是用ldd查看 TensorRT 库依赖了哪个 CUDA、哪个 cuDNN,然后通过 LD_LIBRARY_PATH 的先后顺序强制它选对版本。

3. 模型转换:从 PyTorch 到 ONNX 再到 TensorRT

3.1 PyTorch 导出 ONNX 时的关键参数

前面铺垫了这么多,现在进入正题。整个转换链路通常是 PyTorch -> ONNX -> TensorRT,第一步是 PyTorch 导出 ONNX。这个环节看似简单,但很多参数设置不当会直接影响后续 TensorRT 引擎的质量。

我拿 YOLOv5s 举例,导出命令大概是这样的:

import torch from models.experimental import attempt_load # 加载训练好的权重 model = attempt_load('weights/best.pt', map_location='cpu') model.eval() # 构造一个固定尺寸的输入 dummy_input = torch.zeros(1, 3, 640, 640) # 导出 ONNX torch.onnx.export( model, dummy_input, 'yolov5s.onnx', opset_version=11, input_names=['images'], output_names=['output'], dynamic_axes={ 'images': {0: 'batch_size'}, 'output': {0: 'batch_size'} } )

几个关键点的说明:

opset_version 的选择。TensorRT 8.x 对 ONNX opset 11 的支持最成熟,opset 13 开始有一些算子的表述方式变了,TensorRT 的解析器不一定能完美支持。我之前用过 opset 13 导出,结果在 TensorRT 转换时报了一个莫名其妙的 "Unsupported layer" 错误,换成 opset 11 之后顺利通过。如果你用的模型里有比较新的算子导致 opset 11 导出失败,那可能要检查模型结构,而不是强上高版本 opset。

dynamic_axes 的设定。如果你的部署场景要求 batch 可变,或者输入分辨率可变,这里必须设置 dynamic_axes。但要注意,TensorRT 处理动态 shape 需要额外做 shape 优化(后面会详细说),如果你场景固定,比如永远 640x640、一次一张图,建议把动态轴去掉,能获得更极致的性能。

模型里不要带 DDP 包装。很多人训练时用 DistributedDataParallel 包装了模型,直接导出会带上一堆冗余结构,需要用model.module解包后再导。还有个容易忽略的点:导出前一定要确认模型处于 eval 模式,并且把 BatchNorm 层的 running_mean 和 running_var 冻结,否则导出的 ONNX 里 BN 层行为不对,推理结果会偏离预期。

导出完成后可以用onnx.checker和可视化工具(比如 Netron)检查一遍模型的输入输出节点是不是跟预期一致。我曾经导出过一个语义分割模型,前向输出是两三个不同尺度的特征图,我只取了一个导出,导致后续整个检测头的结果都不对,这个在可视化阶段就能发现。

3.2 用 trtexec 还是 Python API 做转换

拿到 ONNX 模型之后,转换成 TensorRT 引擎有两条路:命令行工具 trtexec 和 Python 脚本。

trtexec 的优势是快、简单、参数直白,适合快速验证模型能否被正确解析、性能和精度大概什么水平。我强烈建议你第一次转换先用 trtexec 跑通,再去写 Python API 的封装代码。

# FP32 引擎 /opt/TensorRT-8.2.3.0/bin/trtexec \ --onnx=yolov5s.onnx \ --saveEngine=yolov5s_fp32.engine \ --workspace=2048 # FP16 引擎 /opt/TensorRT-8.2.3.0/bin/trtexec \ --onnx=yolov5s.onnx \ --saveEngine=yolov5s_fp16.engine \ --workspace=2048 \ --fp16

--workspace参数指定 TensorRT 在构建引擎时可以使用的显存上限(单位 MB)。这个值设太小会导致某些融合层无法构建,设太大对性能没有额外提升但会占用显存。我一般设成 2048 或 4096,你要是在显存紧张的设备上跑,可以试试 1024,必要的时候再逐步往上加。

用 trtexec 转换完,除了.engine文件,命令行还会输出详细的性能报告,里面有 min/mean/max 延迟、吞吐量这些数据。我在做优化前会先跑一遍 trtexec 基线,有了一个客观的起点,后面调优就能判断方向是否正确。

Python API 转换的代码稍长一些,但更灵活,可以自定义优化选项、推理时的显存策略、动态 shape 的优化范围。我常用的转换函数模板:

import tensorrt as trt def build_engine(onnx_path, engine_path, precision='fp16', dynamic=False, min_shape=(1,3,640,640), opt_shape=(4,3,640,640), max_shape=(8,3,640,640)): 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(onnx_path, 'rb') as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(parser.get_error(i)) return None config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30) # 2GB if precision == 'fp16': config.set_flag(trt.BuilderFlag.FP16) if dynamic: profile = builder.create_optimization_profile() input_name = network.get_input(0).name profile.set_shape(input_name, min_shape, opt_shape, max_shape) config.add_optimization_profile(profile) # 创建序列化引擎 serialized_engine = builder.build_serialized_network(network, config) if serialized_engine is None: print("Engine build failed") return None with open(engine_path, 'wb') as f: f.write(serialized_engine) print(f"Engine saved to {engine_path}") return engine_path

这个函数里我保留了动态 shape 的优化选项,min/opt/max 三个 shape 的设定非常影响性能。TensorRT 会用 opt_shape 去做 kernel 选择,如果你的实际推理 batch 数和 opt_shape 差太远,性能会大打折扣。比如你实际跑 batch=4,opt 却设了 batch=8,TensorRT 选出来的实现可能在 batch=4 上并不是最优的。我通常会把 opt_shape 设成和线上实际使用最频繁的 batch 一致,偏差超过一倍就需要重新评估。

还有一点:引擎文件的生成和设备绑定。同一个 ONNX 在不同型号的显卡上生成的引擎不能混用,比如你在 2080Ti 上生成的 .engine 文件拿到 T4 上直接加载会报 "Unsupported planner",需要用 T4 重新构建一次。这块在生产环境里要提前规划,建议做成模型构建服务,按设备类型异步生成引擎缓存。

4. 推理性能实测:T4 上的单路与多路对比

4.1 测试方案设计

引擎构建好之后,最关键的环节就是实测对比。我当时的测试目标是验证"从 ONNX Runtime 到 TensorRT 原生,到底能带来多大的提升",所以设计了三个对比项:ONNX Runtime CUDA 执行提供程序、TensorRT FP32、TensorRT FP16,分别在单路和 8 路并发两种场景下测延迟和吞吐量。

测试环境就是前面说的 T4 + CUDA 11.4 + TensorRT 8.2.3,模型是 YOLOv5s,输入分辨率固定 640x640,测试视频流是 1080p@25fps 的交通监控画面。延迟指标我取的是端到端延迟,包括图像预处理(缩放、归一化、通道转换)、推理、后处理(NMS)三个环节的总耗时;吞吐量指标用"每秒处理多少张图"来算,多路场景模拟的是 8 个独立视频流且每路独立输入的情况。

特别说明一下,预处理和后处理我全程用 CUDA 完成,没有把数据从 GPU 拷贝回 CPU 再处理,因为一旦数据传输成为瓶颈,引擎再快也没用。具体做法是用 Python 的 cupy 或 torch 的 CUDA tensor 来做归一化和缩放,NMS 我用的 TensorRT 提供的 EfficientNMS 插件版本,而不是传统 PyTorch 实现的 NMS,这一步能省掉不少耗时。

4.2 实测数据整理与对比

我整理了一份比较典型的实测数据(因为批次和数据集的差异,你的环境跑出来可能会有些波动,但整体趋势应该是一致的):

方案单路延迟 (ms)8路并发延迟 (ms/路)8路总吞吐量 (fps)GPU 利用率
ONNX Runtime CUDA12.126.8预设约 13042%
TensorRT FP329.419.2约 17568%
TensorRT FP166.212.5约 21589%

这个表格里 ONNX Runtime CUDA 的数字是我自己实测的,8 路并发时需要 26.8ms/路,意味着单卡最多撑不到 5 路流畅 1080p@25fps 的场景,而 TensorRT FP16 的 12.5ms/路,理论上可以支撑 6 到 7 路,实测把视频流降到 20fps 帧率后可以跑到 8 路。这也是为什么我标题里写"1080p25帧用 TensorRT YOLO 640 分辨率能支持多少路",我的答案是:单张 T4 在 FP16 + batch 并发的前提下,支持 6-8 路是靠谱的。

再看 GPU 利用率,ONNX Runtime 那个 42% 很能说明问题。它并不是算不快,而是很多时间浪费在等待数据加载和内核启动上,GPU 的空闲率太高。TensorRT 的层融合和内核自动调优把计算密度拉满,利用率自然就上去了。

我后来还尝试了把 8 路输入合并成 batch=8 一起推理,而不是每路单独推理。在不改变模型结构的前提下,batch=8 的总吞吐量比 8 路独立推理高了约 18%,因为 GPU 的 Tensor Core 更擅长大批量矩阵乘法。但代价是延迟会略微上升,因为要等待 8 帧数据全部到位才能开始推理。如果你的系统对单帧延迟敏感,就不建议用大 batch;如果是做离线批量视频分析,可以用大 batch 拿吞吐量。

4.3 性能调优的进阶手段:多流与流水线

实测过程中我还发现了一个很容易被忽略的性能瓶颈:GPU 的利用率高不代表你的代码没有浪费。如果你推理完一个 batch 之后要等后续视频帧到达,GPU 就处于空闲状态。这个问题的标准解法是 CUDA Stream 多流并行。

TensorRT 的推理接口支持指定 CUDA Stream,你可以把不同视频流的预处理、推理、后处理放到不同的 stream 上。T4 有多个 copy engine 和计算引擎,多流可以让数据搬运和计算重叠起来。我实测下来,把 8 路视频帧均匀分配到 4 个 stream 上,整体吞吐量能再涨 10-15%。

stream = torch.cuda.Stream() with torch.cuda.stream(stream): # 在指定 stream 上执行预处理 processed = preprocess(frames) # 让 TensorRT 上下文使用这个 stream context = engine.create_execution_context() context.set_optimization_profile_async(0, stream.cuda_stream)

这块代码的细节有点多,你如果没接触过 CUDA Stream 的概念,可以先记住一个结论:传入 TensorRT 执行接口的所有输入数据、输出结果的分配和计算都要挂在同一个 stream 上,并且这个 stream 要和 PyTorch 默认的 stream 分开,否则会出现数据不同步的脏读脏写问题,表现是偶发性的预测结果漂移,排查起来非常费劲。

5. 常见问题与排查技巧实录

5.1 版本不匹配的一类经典报错

我前面说版本问题时提过一个偶发情况,这里展开讲。TensorRT 加载引擎文件时报 "Deserialize the cuda engine failed" 是特别常见的,查下来十有八九是两种情况:一是 .engine 文件是从别的显卡型号构建的,跟前面对不上;二是 TensorRT 运行时和构建时的依赖库版本变了,比如 cudnn 从 8.2 升级到了 8.4,旧引擎反序列化失败。

第二种情况比较隐蔽,因为你明明没有改代码,就是把环境里的 cuDNN 更新了一下,就全崩了。这背后的原因:TensorRT 引擎文件里保存了构建时的优化信息,其中包括某些算子的实现细节,这些细节是绑定到具体库版本的。我的建议是生产环境里把 TensorRT、CUDA、cuDNN 的版本全部固定,用条件写进项目文档里,不要随便升级任何一个组件。这也算是部署工程师的基本素养,只升级你要升级的东西,不顺手升级整个环境。

还有一个值得提醒的点:TensorRT 的 Python API 和 C++ API 版本不一致也会导致类似问题。我自己遇到过用 Python 构建的引擎,C++ 程序加载时报版本校验错误,最后发现是 Python 环境的 tensorrt 包和 /opt/TensorRT 下的 C++ 库不是同一个版本。所以在同一台机器上做 Python 环形测试和 C++ 生产部署时,要确保所有环境变量、pip 安装包的来源是同一份 TensorRT 安装包。

5.2 动态 batch 与动态分辨率处理

如果你用了 dynamic_axes 导出 ONNX,在 TensorRT 转换时也设置了 optimization_profile,那推理阶段就需要在每次推理前设置输入/输出的实际 shape。这个环节有个新手容易踩的坑:输入张量的内存分配大小必须按 max_shape 来分配,即使这次实际只跑 batch=1。

# 假设 max_shape = (8, 3, 640, 640) input_shape = (1, 3, 640, 640) context.set_input_shape('images', input_shape) # 注意:d_input 是按 max_shape 分配的 d_input = torch.zeros((8, 3, 640, 640), dtype=torch.float32, device='cuda') d_output_h = torch.zeros((8, 25200, 6), dtype=torch.float32, device='cuda')

这个分配策略背后的原因是 TensorRT 在引擎构建时就已经根据 max_shape 规划好了显存空间,如果你用小 shape 分配一次,下次换大 shape 时会因为显存不够直接报 out of memory。我之前踩过这个坑,现象是一次推理用完 batch=6,后续跳回 batch=1 的时候报错显存重复占用,实际显存明明还很充足。

关于动态分辨率的处理就更讲究了。如果你的模型结构里有下采样倍数(比如 YOLO 是 32 倍下采样),那么输入宽度和高度必须是 32 的倍数。我试过 640x640 稳定运行,改成 642x642 之后检测框完全乱套,就是因为特征图尺寸对不齐。所以做动态分辨率的时候,一定要在预处理阶段把实际输入尺寸调整到对齐值。

5.3 算子不支持与插件缺失

ONNX 模型里偶尔会有 TensorRT 解析器不支持的算子,报错是 "Unsupported Layer" 后面跟着算子名。遇到这种情况先别慌,有几个排查方向。

第一,检查是不是 opset 版本问题导致的。有些算子在低版本 opset 里有对应的解析,高版本反而没有,或者反过来了。最常见的做法是导出时固定 opset=11,再用 trtexec 试一遍,能跳过很多解析问题。

第二,用网络结构里一些不常见的自定义算子,比如我在一个文本检测模型里遇到过自定义的旋转框回归层,TensorRT 原生没有对应实现。解决办法是用 TensorRT 的 plugin API 自己实现一个插件,注册进去。这个流程有点复杂:需要实现 getPluginType、getOutputDimensions、enqueue 等一堆接口,编译成动态库,然后在转换脚本里用parser.register_plugin()注册。我建议第一次做插件时从 GitHub 上找现成的第三方插件库作为模板,比如一些开源企业已经实现好了常见的目标检测后处理算子,直接拿来改比从零写快得多。

第三,实在不想写插件,还有一个退路:把模型拆成两段,不支持算子之前和之后的网络各转一个引擎,中间用 Python/C++ 代码手动做数据拼接和流转。这种方式灵活性和性能都不如单一引擎,但在有些边缘场景里是唯一能落地的方案。

5.4 显存占用异常与推理抖动问题

最后再说一个我调优时经常遇到的问题:显存占用异常飙升。TensorRT 构建引擎时设置的 workspace 大小只影响构建期,不影响运行期显存。运行期显存由引擎本身决定,正常情况下一份 YOLOv5s FP16 引擎运行显存占用大概在 2-3GB 左右,如果看到 6GB 以上异常,通常是构建时显存分配策略不对。

我遇到过最典型的异常是有一次用了较大的 workspace(8GB),结果在显存只有 16GB 的显卡上一跑多路就 OOM,后来把 workspace 降到 2GB 就正常了。这里有一个容易混淆的点:workspace 大的时候 TensorRT 更倾向于使用"占用大但执行快"的内核,如果你的显卡显存偏小,就得在工作效率和显存占用之间做平衡。所以我一般建议显存小于 16GB 的设备,workspace 控制在 2-4GB;大显存设备(比如 24GB 的 3090)可以放到 6GB 以上去追求极致性能。

还有一种推理抖动情况:同一份引擎每次运行延迟波动很大,一会 6ms 一会 18ms。这个大概率是 CPU 和 GPU 之间的同步问题造成的,比如推理前没有把输入数据从 CPU 侧同步到 GPU,导致 TensorRT 启动时隐式做了一个阻塞式拷贝。解决方式是把数据拷贝的代码显式放到 TensorRT 调用之前,并记录到日志里,排查数据是否会成为瓶颈。另一个隐藏原因是显存碎片化,多路推理中频繁分配释放显存会让显存分配器产生碎片,定期用固定批量推理可以减少这种情况。

6. 最后想分享的一点心得

整个从 ONNX Runtime 到 TensorRT 原生的过程走完,我最大的感悟是:推理加速不是简单地换个引擎就算完事,它是一个系统性地优化输入输出链路、显存调度方式、并发模式和精度策略的过程。你换到 TensorRT 之后如果只是把模型文件替换了,数据预处理后处理还是走 CPU,那提升幅度可能只有 20-30%;但如果你把预处理的归一化放到 GPU 上、把后处理 NMS 也做成 GPU 算子、用多流调度把数据拷贝和计算重叠,性能提升才能倍数级显现。

还有一个很实用的小技巧想分享给你:换引擎之后一定要自己验证一遍精度。你可以把同一批输入分别用 PyTorch 原模型和 TensorRT 引擎推理,计算二者的输出差,FP16 下差值一般在 0.01 量级,INT8 会更大。我习惯写一个自动阈值脚本,超过阈值就抛错,这样每次构建新的引擎都能快速验证是否满足要求。这个习惯在你后续做模型更新、训练数据迭代的时候特别有用,不然模型一换,引擎精度掉点,线上出问题的时候你压根不知道是训练的问题、导出的问题还是 TensorRT 转换的问题。

如果你正卡在某一步,建议按这个顺序排查:先确认环境版本,再跑一遍 trtexec,确认模型解析没问题,再写 Python 封装,逐步加动态 shape 和多流逻辑。每一步都验证通过再进入下一步,你会发现问题越来越少,对这套流程的理解也会越来越深。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询