Polygraphy 与 TensorRT 深度互操作:利用 extend 装饰器在保留原生 TensorRT API 的同时复用 Polygraphy 加载器 📅 发布时间:2026/9/15 12:24:08 👁 浏览次数: Polygraphy 与 TensorRT 深度互操作利用 extend 装饰器在保留原生 TensorRT API 的同时复用 Polygraphy 加载器【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT导读本篇文章以 Polygraphy 官方示例03_interoperating_with_tensorrt为核心讲解如何在不放弃 Polygraphy 便捷加载器Loader的前提下随时切入原生 TensorRT API 完成网络修改、Builder 配置等高级操作。读完本文你将掌握func.extend()装饰器的完整语义、NetworkFromOnnxPath与CreateConfig的内部返回结构以及如何通过先扩展、后接管的组合方式在构建引擎前修改网络、设置 Polygraphy 暂未封装或你希望手动控制的 TensorRT Builder Flag。该示例位于仓库 tools/Polygraphy/examples/api/03_interoperating_with_tensorrt 目录下。互操作设计理念Polygraphy 不隐藏后端 APIPolygraphy 的关键设计原则是与 TensorRT 及其他后端完全互操作它不会隐藏底层后端 API因此你可以在 Polygraphy API 与后端 API例如 TensorRT之间自由切换。这意味着当你需要 TensorRT 提供的高级能力时可以直接拿到trt.Builder、trt.INetworkDefinition、trt.OnnxParser、trt.IBuilderConfig等原生对象进行操作同时继续享受 Polygraphy 在加载、构建、运行与调试上的便利——也就是官方文档所说的 the best of both worlds两全其美。本示例聚焦于两个最常见的互操作场景在构建引擎之前修改 TensorRT 网络例如设置网络名称、增减层、改写张量精度设置 Polygraphy 尚未支持或你希望手动精确控制的 TensorRT Builder Flag。示例的核心入口是 example.py依赖仅有numpy见同目录下 requirements.txt模型文件为identity.onnx一个恒等映射模型输入x、输出y。运行示例三步复现按照官方文档运行该示例只需三步安装前置依赖确保 TensorRT 已正确安装Python 侧为tensorrt包安装其他依赖python3 -m pip install -r requirements.txt即numpy。可选检查load_network()生成的 TensorRT 网络load_network是示例脚本中定义的函数下面的命令会从脚本内部调用它并把生成的网络以名称应为MyIdentity的形式展示出来polygraphy inspect model example.py --trt-network-func load_network --show layers attrs weights其中--trt-network-func指示 Polygraphy 直接调用脚本中自定义的、返回 TensorRT 网络的函数--show layers attrs weights控制展示内容网络层、层属性与权重是排查网络结构时的常用组合。运行示例python3 example.py预期输出包含Network name: MyIdentity与Inference succeeded!并最终通过断言np.array_equal(outputs[y], inp_data)——因为这是一个恒等模型输出应当与输入完全一致。核心机制一用func.extend()扩展NetworkFromOnnxPath示例中第一个装饰用法如下from polygraphy.backend.trt import ( CreateConfig, EngineFromNetwork, NetworkFromOnnxPath, TrtRunner, ) from polygraphy import func # 被扩展的加载器NetworkFromOnnxPath(identity.onnx) func.extend(NetworkFromOnnxPath(identity.onnx)) def load_network(builder, network, parser): # 在这里可以任意修改网络 network.name MyIdentity print(fNetwork name: {network.name}) # 注意无需 return 任何值extend() 会替你转发返回值要点在于被装饰函数的参数个数与类型必须与被扩展加载器的返回值一一对应。根据 loader.py 中NetworkFromOnnxPath.call_impl()的实现返回语句为return builder, network, parserNetworkFromOnnxPath会返回三元组(trt.IBuilder, trt.INetworkDefinition, trt.OnnxParser)——这正是load_network(builder, network, parser)收到的三个参数。从源码看NetworkFromOnnxPath的解析流程是先调用create_network()创建空网络随后trt.init_libnvinfer_plugins()初始化插件库、trt.OnnxParser(network, logger)创建解析器最后parser.parse_from_file(path)从文件解析模型——选用parse_from_file是为了让解析器能够跟踪 ONNX 文件位置以便处理外部权重external weights。此外NetworkFromOnnxPath还支持若干可选参数见 loader.py 构造函数flagstrt.OnnxParserFlag列表用于修改解析器默认行为plugin_instancenorm置True时强制使用 InstanceNorm 的插件实现清除NATIVE_INSTANCENORM标志strongly_typed是否将网络标记为强类型strongly typed。NetworkFromOnnxBytes与之等价只是改为从内存字节解析模型。核心机制二用func.extend()扩展CreateConfig设置 Builder Flag第二个装饰用法针对构建配置func.extend(CreateConfig()) def load_config(config): # Polygraphy 本身支持 fp16 标志但假如它不支持我们可以这样手动设置 config.set_flag(trt.BuilderFlag.FP16)由于CreateConfig的返回值是trt.IBuilderConfig因此load_config的形参config直接就是原生 TensorRT 的 builder 配置对象可以调用config.set_flag(trt.BuilderFlag.FP16)这类标准 TensorRT API。需要说明的是Polygraphy 的CreateConfig其实已经封装了大量常用配置项。从 config.py 的构造函数可见其显式参数包括参数含义默认值tf32是否启用 TF32 精度Falsefp16是否启用 FP16 精度Falsebf16是否启用 BF16 精度Falseint8是否启用 INT8 精度Falsefp8是否启用 FP8 精度FalsecalibratorINT8 校准器trt.IInt8Calibrator网络无显式精度且使用 INT8 时需要Noneuse_dla[实验性] 是否将 DLA 设为默认设备Falseallow_gpu_fallback[实验性] DLA 开启时是否允许层回退到 GPUFalse其余配置通过**kwargs透传给基类_CreateConfigCommon包括动态形状的profiles、precision_constraintsobey/prefer、restricted对应SAFETY_SCOPE、refittableREFIT、strip_planSTRIP_PLAN、direct_ioDIRECT_IO、sparse_weightsSPARSE_WEIGHTS、memory_pool_limits内存池上限、tactic_sources、timing_cache_path策略时序缓存路径、algorithm_selector、engine_capability等见 config.py 的配置应用逻辑。extend的价值正在于此当某次构建需要用到 Polygraphy 尚未封装、或者你希望绕过其默认行为的 Builder Flag例如组合set_flag与set_memory_pool_limit等底层调用时不必放弃 Polygraphy 的加载链路只需在装饰函数里用原生 API 补齐即可。extend装饰器的完整语义源码级解读extend定义在 tools/Polygraphy/polygraphy/func/func.pypolygraphy.func.extend其本质是被装饰函数y以被扩展函数x的返回值为参数运行然后把x的返回值继续向前传递。具体规则返回值自动转发若y没有return或返回Noneextend会把x的返回值原样返回给调用者因此y对外提供与x完全一致的接口——示例中的load_network不写return正是利用这一点覆盖语义若y返回非None的值则该值会取代x的返回值交给调用者参数透传若需要同时访问x的原始入参可以让y在形参列表前面加上x的参数此时所有传给x的参数都会被转发给y。注意若x原地修改了参数y看到的是修改后的对象自动解包元组x返回的元组会被自动解包y按解包后的顺序接收参数限制被装饰函数不能使用*args/**kwargs变长参数参数个数不匹配时G_LOGGER.critical会直接报错并提示期望的参数列表。上述行为与示例注释我们不需要返回任何东西——extend()会替我们处理完全吻合。组合装配从加载器回到常规 Polygraphy 流程在完成网络与配置的扩展之后示例回到常规的 Polygraphy API 完成构建与推理def main(): # 使用懒加载lazy加载器时传入的是函数本身而不是调用它们 build_engine EngineFromNetwork(load_network, configload_config) with TrtRunner(build_engine) as runner: inp_data np.ones(shape(1, 1, 2, 2), dtypenp.float32) # 注意runner 拥有输出缓冲区并在多次 infer() 之间复用 # 如需保存多次推理结果请使用 copy.deepcopy() outputs runner.infer({x: inp_data}) assert np.array_equal(outputs[y], inp_data) # 恒等模型输出应等于输入 print(Inference succeeded!)几点值得展开EngineFromNetwork继承自EngineBytesFromNetwork见 loader.pynetwork参数既可以是(builder, network[, parser])元组也可以是返回该元组的可调用对象——这里传入load_network函数本身懒加载。构建时EngineFromNetwork会调用传入的config可调用对象拿到trt.IBuilderConfig尝试挂载 Polygraphy 校准器然后调用builder.build_serialized_network(network, config)完成引擎构建旧版本回退到build_engineserialize并可选地在timing_cache_path指定的路径保存/合并策略时序缓存。配置阶段会先创建空时序缓存config.create_timing_cache(b)确保构建过程中可以被填充。懒加载与立即求值EngineFromNetwork(load_network, configload_config)中传的是函数引用而非调用结果。这是 Polygraphy 的懒加载lazy风格与之对应的立即求值函数式 APIimmediate evaluation functional API会让互操作更加直白示例注释建议参考 examples/api/06_immediate_eval_api对应examples/api/06_immediate_eval_api目录。TrtRunner见 runner.py接受引擎或可返回引擎的可调用对象激活时自动创建执行上下文、分配输入/输出缓冲区与 CUDA 流infer()将输入字典送入设备并返回输出字典。其构造参数还包括optimization_profile多优化档引擎要激活的 profile 索引默认使用第 0 个 profile可通过set_profile()动态切换allocation_strategy执行上下文设备内存分配策略取值为static默认按所有 profile 的最大形状预分配、profile按当前 profile 的最大形状、runtime按当前输入形状后两者使用trt.ExecutionContextAllocationStrategy.USER_MANAGEDweight_streaming_budget/weight_streaming_percent运行时权重流式weight streaming的显存预算与驻留比例控制。需要留意的是runner 被官方定位为面向原型验证、测试与调试并不建议直接用于生产部署。常见问题与排查建议polygraphy inspect model example.py --trt-network-func load_network找不到函数确认脚本中该函数存在且位于模块顶层函数名需与--trt-network-func完全一致。断言np.array_equal(outputs[y], inp_data)失败恒等模型理论上输出等于输入若失败应优先检查identity.onnx是否被替换、输入张量名是否为x、输出是否为y。extend报参数数量不匹配根据G_LOGGER.critical输出的期望参数列表核对被扩展加载器的真实返回值例如扩展NetworkFromOnnxBytes/NetworkFromOnnxPath需 3 个参数而扩展仅返回(builder, network)的加载器如CreateNetwork只需 2 个。想要更系统地查看网络polygraphy inspect model支持--show layers attrs weights组合可分别查看网络层、层属性和权重明细是理解 ONNX 解析结果与验证自定义修改是否生效的利器。小结03_interoperating_with_tensorrt示例展示了 Polygraphy 互操作能力的正确打开方式利用func.extend()在加载链路上插入自定义逻辑直接操作原生 TensorRT 对象网络、builder 配置再无缝交还给 Polygraphy 完成引擎构建与推理。无论是修改网络、设置额外 Builder Flag还是接入自定义校准器与算法选择器这一模式都能让你在享受 Polygraphy 便利性的同时完整保留 TensorRT 底层 API 的全部能力。深入阅读 func.py、loader.py、config.py 与 runner.py 的实现可以进一步理解加载器返回值的契约以及配置项在构建期的实际作用为自定义工作流打下坚实基础。【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考