TensorRT 使用技巧:用 Python API 与 trtexec 打通 ONNX 到 plugin 的配置骨架
发布时间:2026/9/26 9:14:35来源:尧图网络
1. 从 ONNX 到 engine一条容易踩坑的部署链路TensorRT 模型部署这件事说简单也简单说麻烦也麻烦。简单在于核心流程就三步ONNX 导出、trtexec 验证、Python API 加载 engine。麻烦在于每一步都有细节能把人卡住——ONNX 的 opset 版本不对、动态 shape 没配好、plugin 没注册导致 engine 反序列化直接报错、CUDA_VISIBLE_DEVICES 设置晚了导致跑到了错误的卡上。这篇聚焦的是「配置骨架」给你一套可以直接复制、逐项验证的 trtexec 参数、Python 推理脚本和 plugin 注册片段。适合已经在本地或服务器上装好 TensorRT、手里有一个 ONNX 文件、想快速跑通推理链路的人。如果你还没装 TensorRT建议先用官方 tar 包或 pip 装好 tensorrt 和 pycuda再往下看。我试过在同一个模型上反复折腾 ONNX 导出参数和 trtexec 的 shape 配置最后发现大部分精度和性能问题都出在「导出时没对齐」和「构建时没指定」这两件事上。下面按链路顺序拆开讲。2. TaoToken 前置把模型对话和 API Key 准备好在开始写 TensorRT 脚本之前有一个容易被忽略的前置动作确认你手头的模型文件来源和版本信息。很多时候 ONNX 是从某个对话模型或代码模型导出的导出参数、输入输出名字、opset 版本都需要对齐。这时候可以用 TaoToken 的模型对话能力快速确认模型结构描述或者用它来生成一段 ONNX 导出脚本的骨架。TaoToken 的入口在这里官网https://taotoken.net/?utm_sourcetaotoken_aicg_blog_endutm_mediumcsdnutm_campaignrewriteutm_contentAPIhttps://taotoken.net/api模型对话https://taotoken.net/models?utm_sourcetaotoken_aicg_blog_endutm_contentmodel_chatutm_campaignrewriteAPI Keyshttps://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite接入文档https://taotoken.net/doc?utm_sourcetaotoken_aicg_blog_endutm_contentdocutm_campaignrewrite如果你后面要长期做编码和 Agent 相关的模型部署可以关注 Coding Planhttps://taotoken.net/coding-plan?utm_sourcetaotoken_aicg_blog_endutm_contentcoding_planutm_campaignrewrite拿到 API Key 之后你可以用它来辅助生成 trtexec 的参数模板或者让模型帮你检查 ONNX 的输入输出维度是否和 engine 构建时一致。这一步不是必须的但能省掉不少「为什么 trtexec 报 shape mismatch」的排查时间。3. 可复制配置trtexec 参数与 Python API 骨架3.1 ONNX 导出时的对齐检查在跑 trtexec 之前先确认 ONNX 的输入输出名字和维度。用下面这段脚本打印出来import onnx model onnx.load(model.onnx) print(opset:, model.opset_import) for inp in model.graph.input: dims [d.dim_value if d.dim_value 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 else d.dim_param for d in out.type.tensor_type.shape.dim] print(output:, out.name, dims)如果看到dim_param是空的或者batch这种字符串说明是动态 shapetrtexec 构建时需要显式指定--minShapes、--optShapes、--maxShapes。如果全是dim_value那就是静态 shape直接构建即可。3.2 trtexec 构建 engine 的完整参数静态 shape 的构建命令trtexec \ --onnxmodel.onnx \ --saveEnginemodel.engine \ --fp16 \ --workspace4096 \ --verbose \ --dumpProfile \ --exportProfileprofile.json动态 shape 的构建命令假设输入名叫input维度是[batch, 3, 224, 224]trtexec \ --onnxmodel.onnx \ --saveEnginemodel_dynamic.engine \ --minShapesinput:1x3x224x224 \ --optShapesinput:8x3x224x224 \ --maxShapesinput:16x3x224x224 \ --fp16 \ --workspace4096 \ --verbose几个参数的含义对照参数作用建议值--fp16启用 FP16 精度精度允许时开启--workspace构建时可用显存 MB4096 起步--dumpProfile打印每层耗时排查性能瓶颈--exportProfile导出 profile JSON配合可视化--minShapes动态 shape 最小值按实际 batch 下限--optShapes动态 shape 最优值常用 batch--maxShapes动态 shape 最大值显存上限注意--workspace在 TensorRT 8.x 之后语义有变化部分版本用--memPoolSizeworkspace:4096替代。如果报未知参数先trtexec --help确认当前版本支持的写法。3.3 Python API 加载 engine 并推理构建好 engine 之后用 Python API 加载并跑一次推理import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np TRT_LOGGER trt.Logger(trt.Logger.WARNING) def load_engine(engine_path): with open(engine_path, rb) as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def infer(engine, input_data): with engine.create_execution_context() as context: input_name engine.get_tensor_name(0) output_name engine.get_tensor_name(1) context.set_input_shape(input_name, input_data.shape) output_shape context.get_tensor_shape(output_name) output_data np.empty(output_shape, dtypenp.float32) d_input cuda.mem_alloc(input_data.nbytes) d_output cuda.mem_alloc(output_data.nbytes) stream cuda.Stream() cuda.memcpy_htod_async(d_input, input_data, stream) context.set_tensor_address(input_name, int(d_input)) context.set_tensor_address(output_name, int(d_output)) context.execute_async_v3(stream_handlestream.handle) cuda.memcpy_dtoh_async(output_data, d_output, stream) stream.synchronize() return output_data engine load_engine(model.engine) x np.random.randn(1, 3, 224, 224).astype(np.float32) y infer(engine, x) print(output shape:, y.shape)这段脚本的关键点在于set_tensor_address和execute_async_v3的配合。TensorRT 8.5 之后推荐用execute_async_v3旧的execute_async在部分版本已经废弃。3.4 plugin 注册的 config 骨架如果你用了自定义 pluginengine 反序列化时会去找对应的 plugin creator。注册骨架如下import tensorrt as trt class MyPluginCreator(trt.IPluginCreator): def __init__(self): super().__init__() self.name MyPlugin self.version 1 self.plugin_namespace self.field_names trt.PluginFieldCollection([ trt.PluginField(scale, np.array([1.0], dtypenp.float32), trt.PluginFieldType.FLOAT32), ]) def create_plugin(self, name, field_collection, plugin_namespace): scale field_collection[0].data[0] return MyPlugin(scale) def register_plugin(): registry trt.get_plugin_registry() creator MyPluginCreator() registry.register_creator(creator, )注册之后再调用load_engine就不会报Plugin not found了。注意plugin_namespace要和导出 ONNX 时写入的 namespace 一致否则匹配不上。4. 验证请求与成功结果4.1 trtexec 验证构建完成后trtexec 会输出一段性能摘要。关注这几个指标trtexec --loadEnginemodel.engine --shapesinput:1x3x224x224 --iterations100 --avgRuns10输出里会有Throughput和Latency两栏。如果Throughput是 0 或者报错先检查--shapes是否和 engine 构建时的动态范围匹配。4.2 Python 推理验证跑上面那段 Python 脚本如果输出 shape 和预期一致说明链路通了。再对比一下 ONNX Runtime 的输出import onnxruntime as ort sess ort.InferenceSession(model.onnx) ort_out sess.run(None, {input: x})[0] print(max diff:, np.max(np.abs(ort_out - y)))max diff在 1e-3 以内算正常FP16 下可能到 1e-2。如果差得离谱优先查输入预处理是否一致。4.3 plugin 验证如果用了 plugin在load_engine之前打印一下已注册的 creatorregistry trt.get_plugin_registry() for c in registry.get_creator_list(trt.TensorRTPhase.BUILD): print(c.name, c.version)确认你的 plugin 名字在列表里再加载 engine。5. 本篇常见错排查5.1CUDA_VISIBLE_DEVICES设置无效这个坑很典型。os.environ[CUDA_VISIBLE_DEVICES]2必须放在import pycuda.autoinit之前否则 pycuda 初始化时已经枚举了所有卡设置就晚了。import os os.environ[CUDA_VISIBLE_DEVICES] 2 import pycuda.autoinit import tensorrt as trt5.2Plugin not found或Serialization failedengine 是在构建时把 plugin 信息序列化进去的反序列化时如果找不到对应的 creator就会直接失败。排查顺序先确认 plugin 名字和版本号一致再确认 namespace 一致最后确认register_creator在deserialize_cuda_engine之前调用。5.3 动态 shape 报shape mismatchtrtexec 构建时用了--minShapes/--optShapes/--maxShapes推理时传入的 shape 必须落在这个范围内。如果 Python 脚本里set_input_shape传了一个超出范围的维度execute_async_v3会报错。建议在set_input_shape之后打印context.get_tensor_shape(input_name)确认。5.4 精度掉得厉害先确认 FP16 是不是罪魁祸首用--fp16构建一个 engine再用--best构建一个对比输出。如果 FP16 掉点严重考虑对敏感层保持 FP32或者用--precisionConstraints做混合精度。另外INT8 校准需要校准集校准集分布和真实数据差太多也会掉点。5.5pycuda内存分配失败cuda.mem_alloc失败通常是显存不够。检查--workspace是不是设太大了或者 batch 是不是超了。可以用nvidia-smi看当前显存占用再调整--maxShapes。6. 继续把链路跑稳TensorRT 部署这条链路核心就是「对齐」两个字ONNX 导出时对齐 opset 和 shapetrtexec 构建时对齐动态范围Python API 加载时对齐 plugin 注册。每一步都有验证动作不要跳过。如果你在排查精度问题时需要快速对比不同模型的输出可以用 TaoToken 的模型对话来生成对比脚本或者分析 profile 数据https://taotoken.net/models?utm_sourcetaotoken_aicg_blog_endutm_contentmodel_chatutm_campaignrewrite接入相关的 API Key 和文档在这里https://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite 和 https://taotoken.net/doc?utm_sourcetaotoken_aicg_blog_endutm_contentdocutm_campaignrewrite最后一个小技巧把 trtexec 的--exportProfile输出的 JSON 丢给可视化工具能直接看到哪一层耗时最长。很多时候性能瓶颈不在卷积而在 reshape 或者 transpose 这类看起来不起眼的操作上。
网站建设高端定制企业官网