ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

TensorRT 实战:用 trtexec 完成 ONNX 模型转换、推理运行与网络性能测试

TensorRT 实战:用 trtexec 完成 ONNX 模型转换、推理运行与网络性能测试 1. 为什么我建议你先用 trtexec 跑通再写代码如果你手里有一个 ONNX 模型想把它部署到 NVIDIA 显卡上做推理第一反应可能是去写 Python 的 TensorRT API。但我自己的习惯是先别急着写代码先用trtexec把整条链路跑一遍。trtexec是 TensorRT 官方自带的一个命令行工具它能做三件事——把 ONNX 转成 TensorRT engine、直接加载 engine 跑推理、以及做吞吐和延迟的性能基准测试。换句话说模型转换、推理运行、网络性能测试这三步它一个命令就能覆盖。它适合谁适合刚接触 TensorRT 部署、想快速验证模型能不能转、转完性能大概多少的人也适合已经在写部署代码、但需要一个独立参照来对比自己代码性能的工程师。因为trtexec内部用的就是 TensorRT 的标准构建和推理流程它跑出来的数字基本就是你代码能拿到的上限参考。我试过在同一个模型上先用trtexec测出延迟和吞吐再去写 Python 推理脚本结果发现脚本慢了 30%最后排查出来是没开 CUDA Graph。如果没有trtexec这个基准我根本不知道自己的代码还有优化空间。所以这篇就按“转换 → 运行 → 性能测试”的顺序把每个环节的参数配置和验证动作讲清楚命令都可以直接复制。在开始之前你需要确认环境里已经装好了 TensorRTtrtexec通常在/usr/src/tensorrt/bin/trtexec这个路径下如果直接敲trtexec提示找不到命令可以把它加到 PATH或者用绝对路径调用。显卡驱动和 CUDA 版本要跟 TensorRT 匹配这个不匹配的话后面会报各种奇怪的错我在第 5 节会专门讲。2. 转换前的准备ONNX 模型检查与 trtexec 环境确认在真正执行转换命令之前有两件事必须先做确认 ONNX 模型本身是合法的以及确认trtexec能正常工作。很多人一上来就敲转换命令结果报错信息看不懂其实问题出在模型导出阶段。先说 ONNX 模型检查。你可以用 Python 的 onnx 包快速看一眼模型的输入输出名称和形状因为后面写--minShapes、--shapes这些参数时输入名字必须和模型里完全一致。执行下面这段import onnx model onnx.load(model.onnx) for inp in model.graph.input: dims [d.dim_value if d.dim_value 0 else d.dim_param for d in inp.type.tensor_type.shape.dim] print(input:, inp.name, dims) for out in model.graph.output: dims [d.dim_value if d.dim_value 0 else d.dim_param for d in out.type.tensor_type.shape.dim] print(output:, out.name, dims)跑完你会看到类似input: images [1, 3, 640, 640]这样的输出。这里的images就是输入张量名后面配置动态形状时要用到。如果维度里出现的是字符串而不是数字说明这个维度是动态的转换时必须用--minShapes、--optShapes、--maxShapes三个参数一起指定范围。接着确认trtexec可用。直接运行trtexec --help如果能看到一大段参数说明说明环境没问题。如果提示 command not found用绝对路径试试/usr/src/tensorrt/bin/trtexec --help能正常输出的话建议做个软链接或者加到 PATH后面命令写起来短一些。另外可以顺手看一下版本trtexec --version版本信息很重要因为不同 TensorRT 版本对 ONNX 的算子支持不一样有些新算子老版本不认转换就会失败。确认完这两点就可以进入转换环节了。这里补充一个我踩过的坑ONNX 模型如果是从 PyTorch 导出的导出时最好把opset_version设成 11 或以上并且用torch.onnx.export时加上dynamic_axes参数来标记动态维度。如果导出时没标记模型里全是静态形状那转换时就不能用动态 batch只能固定 batch 跑。这个在导出阶段就要想清楚后面改起来要重新导出。3. 可复制的 trtexec 转换配置静态与动态 batch 全流程这一节是核心我把静态 batch 和动态 batch 两种转换方式都给出完整命令并且解释每个参数的作用。你可以直接复制把模型路径和输入名换成自己的。先看静态 batch 的转换。假设你的 ONNX 模型输入是固定的1x3x640x640想转成 FP16 精度、workspace 给 1024MBtrtexec --onnxmodel.onnx \ --saveEnginemodel_fp16.engine \ --fp16 \ --workspace1024 \ --explicitBatch \ --verbose这里--onnx指定输入模型--saveEngine指定输出的 engine 文件--fp16开启半精度--workspace是构建时允许使用的显存上限单位 MB。--explicitBatch表示使用显式 batch 维度现在新版本基本都要求加这个。--verbose会打印详细的构建日志第一次转换建议加上方便看每一层的精度选择。如果你想要 TensorRT 自动在所有精度里挑最优的可以用--besttrtexec --onnxmodel.onnx --saveEnginemodel_best.engine --best --workspace1024 --explicitBatch--best会让构建器尝试 FP32、FP16、INT8 等所有可用精度选性能最好的那个。注意 INT8 需要校准数据如果没提供校准集它不会真正启用 INT8只会用 FP16 或 FP32。再看动态 batch 的转换这是实际部署里更常见的场景。假设输入名是images你想支持 batch 从 1 到 8最优 batch 是 4trtexec --onnxmodel.onnx \ --saveEnginemodel_dynamic.engine \ --fp16 \ --workspace2048 \ --minShapesimages:1x3x640x640 \ --optShapesimages:4x3x640x640 \ --maxShapesimages:8x3x640x640 \ --explicitBatch这三个 shapes 参数必须同时设置格式是输入名:batch x 通道 x 高 x 宽。--minShapes是最小形状--optShapes是构建器优化时参考的形状--maxShapes是最大形状。TensorRT 会为optShapes做重点优化所以这个值要设成你实际推理时最常用的 batch。如果你还想在转换后直接跑一次推理验证可以加上--shapes指定实际推理用的形状trtexec --onnxmodel.onnx \ --saveEnginemodel_dynamic.engine \ --fp16 \ --minShapesimages:1x3x640x640 \ --optShapesimages:4x3x640x640 \ --maxShapesimages:8x3x640x640 \ --shapesimages:2x3x640x640 \ --explicitBatch这样转换完会立刻用 batch2 跑一次推理省得你再单独加载 engine。关于 workspace 的设置有个经验值小模型 256MB 够用中等模型 1024MB大模型或者动态形状多的给 2048MB 以上。设太小会导致某些层无法用最优算法性能下降设太大也不会更快只是允许构建器有更多选择空间。我一般先给 1024 试如果日志里出现 “some tactics were not selected due to workspace limit” 之类的提示再往上加。转换完成后你会得到一个.engine文件。这个文件是序列化后的 TensorRT 引擎跟具体的 GPU 型号和 TensorRT 版本绑定换卡或者换版本需要重新转换。这一点要记住不然拿到别的机器上加载会失败。4. 加载 engine 运行推理与性能基准测试实操engine 转好之后下一步就是加载它跑推理并且做性能测试。trtexec加载 engine 用--loadEngine参数性能测试的结果会直接打印在终端里。先看最基本的加载运行trtexec --loadEnginemodel_fp16.engine --batch1--batch指定推理时的 batch size注意这个值不能超过转换时--maxShapes设定的上限。运行完你会看到一段输出里面有Throughput和Latency两个关键指标。Throughput是吞吐量单位是 qps每秒查询数Latency是延迟单位是 ms。这两个是衡量网络性能的核心数据。如果你想让测试更充分可以指定迭代次数和预热次数trtexec --loadEnginemodel_fp16.engine \ --batch4 \ --iterations200 \ --warmUp500 \ --duration10 \ --percentile99--iterations是正式测试的迭代次数--warmUp是预热时间毫秒--duration是至少测试多少秒--percentile99会额外打印 99 分位延迟。这几个参数组合起来测出来的数字比较稳定。我一般会跑--duration10以上太短的测试受抖动影响大。多流测试是提升吞吐量的关键手段。所谓多流就是同时开多个 CUDA stream 并行执行推理适合延迟要求不严格、但吞吐要求高的场景。命令如下trtexec --loadEnginemodel_fp16.engine --batch1 --streams4--streams4表示开 4 个并发流。你可以依次试 2、3、4、8看吞吐量什么时候到瓶颈。通常吞吐会随 stream 数增加而上升但延迟也会变高因为多个流在抢 GPU 资源。你需要根据业务对延迟的容忍度来选。如果想导出时序跟踪信息做深入分析可以加--exportTimestrtexec --loadEnginemodel_fp16.engine --batch1 --exportTimestrace.json生成的trace.json可以用 Chrome 的chrome://tracing打开能看到每个 kernel 的执行时间方便定位瓶颈层。还有一个实用参数是--dumpProfile它会打印每一层的耗时占比trtexec --loadEnginemodel_fp16.engine --batch1 --dumpProfile输出里会按耗时排序排在前面的层就是优化重点。如果发现某个卷积层占了 40% 时间可以考虑对它单独做 INT8 量化或者换更高效的实现。实测下来同一模型 FP32 转 FP16 后延迟通常能降 30% 到 50%吞吐翻倍。如果开了 INT8 且校准得当还能再快一截。但 INT8 精度损失需要自己评估不是所有模型都能接受。5. 常见报错排查从 401 到 OAuth 的对照处理这一节把我遇到过的典型报错和排查思路列出来你对照着看。报错一[E] Error[2]: [builder.cpp::parseModel::xxx] Error Code 2: Internal Error (Assertion ... failed)这种一般是 ONNX 模型里有 TensorRT 不支持的算子。解决办法是先看--verbose日志里具体是哪个算子然后用 ONNX 的onnxsim做简化或者把不支持的算子替换掉。有时候是 opset 版本太高把导出时的 opset 降到 11 或 12 再试。报错二[E] Error[1]: [builder.cpp::buildSerializedNetwork::xxx] Error Code 1: Cuda Runtime (out of memory)显存不够。先降--workspace再降--maxShapes的 batch 上限。如果还不行说明模型本身太大单卡放不下需要考虑模型并行或者换更大显存的卡。报错三[E] Error[3]: [rt.cpp::deserializeCudaEngine::xxx] Error Code 3: API Usage Error (Parameter check failed ...)加载 engine 失败最常见原因是 engine 文件和当前 GPU 或 TensorRT 版本不匹配。engine 是跟硬件绑定的换机器必须重新转换。确认一下转换和加载是不是在同一台机器、同一个 TensorRT 版本下。报错四local proxy failed或connection refused这个通常出现在你通过某种网络方式访问远程服务时。如果你在容器里跑trtexec检查一下网络配置和端口映射。如果是访问外部 API 做模型下载确认网络是通的。这类问题跟trtexec本身无关是环境网络问题。报错五reading choices相关错误这个一般出现在解析模型输出或者配置文件时格式不对。检查你的--shapes参数格式必须是名字:1x3x640x640这种冒号前后不能有空格维度之间用小写字母x连接。报错六OAuth或401 Unauthorized如果你在调用某个云端模型服务或者 API 时遇到 401说明鉴权失败。检查你的 API Key 是否正确、是否过期、请求头里的 Authorization 字段格式对不对。如果是 OAuth 流程确认 token 有没有正确刷新。报错七--explicitBatch相关警告新版本 TensorRT 里--explicitBatch已经是默认行为如果提示这个参数被忽略不用管不影响结果。但如果你的模型是隐式 batch 导出的可能需要加--implicitBatch不过这个模式已经不推荐了。排查的通用思路是先加--verbose看详细日志定位到具体哪一步失败然后单独隔离那一步比如只做转换不推理或者只加载 engine 不转换最后对照 TensorRT 官方文档的 error code 表。大部分问题都能通过日志里的 error code 找到方向。6. 把 trtexec 接入你的部署链路从验证到落地trtexec跑通之后下一步就是把它接入实际的部署链路。这里说几个我常用的做法。第一把trtexec当作 CI 里的模型验证步骤。每次模型更新先跑一遍转换和性能测试确认 engine 能正常生成、延迟和吞吐在预期范围内再进入后续的代码部署。这样可以避免把有问题的模型推到线上。第二用trtexec的输出作为代码优化的基准。你自己写的推理代码性能应该接近trtexec的数字。如果差很多就去检查是不是没开 CUDA Graph、是不是没复用 memory、是不是 stream 配置不对。trtexec内部做了很多优化它的数字就是你的目标。第三动态形状的 engine 在生产里要配合实际请求的 batch 分布来调优。--optShapes设成最常用的 batch能让大多数请求都跑在最优路径上。如果请求 batch 波动很大可以考虑转多个 engine按 batch 范围分别加载。如果你在本地验证完想把这套流程放到远程环境或者团队共享可以用 TaoToken 的 Coding Plan 来管理你的编码任务和配置。它的 API 地址是 https://taotoken.net/api 模型对话入口在 https://taotoken.net/models 接入文档在 https://taotoken.net/doc 。需要生成 API Key 的话去 https://taotoken.net/api-keys 控制台在 https://taotoken.net/console 。这些入口配合你的部署脚本可以把模型转换和性能测试做成自动化流程。最后提醒一点trtexec的 engine 文件不要直接提交到代码仓库它体积大而且跟硬件绑定。正确的做法是把 ONNX 模型和转换脚本提交上去在部署环境里现场转换或者用 CI 在目标机器上生成 engine 再分发。这样换卡换版本时重新跑一遍转换脚本就行不用手动处理 engine 文件。整套流程走下来从 ONNX 到 engine 到性能数字快的话十几分钟就能跑完。先把这条链路跑通再去写部署代码心里会踏实很多。
返回列表