企业官网建设流程全解析

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 不隐藏后端 API

Polygraphy 的关键设计原则是与 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）。

运行示例：三步复现

按照官方文档，运行该示例只需三步：

1. 安装前置依赖：

   • 确保 TensorRT 已正确安装（Python 侧为tensorrt包）；
   • 安装其他依赖：python3 -m pip install -r requirements.txt（即numpy）。

2. （可选）检查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控制展示内容（网络层、层属性与权重），是排查网络结构时的常用组合。

3. 运行示例：

   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(f"Network name: {network.name}") # 注意：无需 return 任何值，extend() 会替你转发返回值

要点在于：被装饰函数的参数个数与类型，必须与被扩展加载器的返回值一一对应。根据 loader.py 中NetworkFromOnnxPath.call_impl()的实现（返回语句为return builder, network, parser），NetworkFromOnnxPath会返回三元组(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 构造函数）：

• flags：trt.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 的构造函数可见，其显式参数包括：

其余配置通过**kwargs透传给基类_CreateConfigCommon，包括：动态形状的profiles、precision_constraints（obey/prefer）、restricted（对应SAFETY_SCOPE）、refittable（REFIT）、strip_plan（STRIP_PLAN）、direct_io（DIRECT_IO）、sparse_weights（SPARSE_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.py（polygraphy.func.extend），其本质是"被装饰函数y以被扩展函数x的返回值为参数运行，然后把x的返回值继续向前传递"。具体规则：

• 返回值自动转发：若y没有return（或返回None），extend会把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, config=load_config) with TrtRunner(build_engine) as runner: inp_data = np.ones(shape=(1, 1, 2, 2), dtype=np.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.py），network参数既可以是(builder, network[, parser])元组，也可以是返回该元组的可调用对象——这里传入load_network函数本身（懒加载）。构建时，EngineFromNetwork会调用传入的config可调用对象拿到trt.IBuilderConfig，尝试挂载 Polygraphy 校准器，然后调用builder.build_serialized_network(network, config)完成引擎构建（旧版本回退到build_engine+serialize），并可选地在timing_cache_path指定的路径保存/合并策略时序缓存。配置阶段会先创建空时序缓存（config.create_timing_cache(b"")），确保构建过程中可以被填充。
• 懒加载与立即求值：EngineFromNetwork(load_network, config=load_config)中传的是函数引用而非调用结果。这是 Polygraphy 的懒加载（lazy）风格；与之对应的"立即求值"函数式 API（immediate 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_MANAGED）；
  • weight_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