企业官网建设流程全解析

使用 Polygraphy CLI 脚本化定义 TensorRT 网络与 Builder 配置（load_network / load_config 实战指南）

【免费下载链接】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

导读

在 TensorRT 推理工作流中，网络结构往往无法直接序列化为磁盘文件，导致命令行工具（如polygraphy run、polygraphy convert）在引擎构建完成之前无法介入。Polygraphy 为此提供了一套约定式脚本接口：只要在 Python 脚本中定义名为load_network（返回 builder、network 与可选 parser）或load_config（返回 builder 配置）的函数，就能把任意用 TensorRT Python API 手工构造或修改的网络直接接入 CLI 工作流。本文将基于 04_defining_a_tensorrt_network_or_config_manually 示例 的完整代码与仓库源码，讲解这套脚本约定的写法、模板生成方法、命令行接入方式及其底层加载机制，读完即可在自己的项目里用一条命令构建并运行自定义网络。

为什么需要"脚本化网络"：CLI 的序列化局限

通常，polygraphy run、polygraphy convert等 CLI 工具的模型参数接受的是 ONNX、TensorFlow 冻结图、engine/plan 等可以被解析或反序列化的文件。但在某些场景下，这种文件模型行不通：

• 你想从零开始用 TensorRT Python API 手工搭建一个网络（例如为了测试某个新算子或算子组合）；
• 你想修改一个由解析器（如 ONNX Parser）创建的网络——例如在解析结果上追加、替换或重排层；
• 这类网络在构建引擎之前无法被序列化到磁盘，也就无法作为模型文件直接传给 CLI。

Polygraphy 的 CLI 工具为此提供了官方绕行方案：如果脚本定义了一个名为load_network的无参函数，返回 TensorRT 的 builder、network 以及可选的 parser，那么就可以把这个脚本文件本身当作模型参数传给 CLI（--model-type=trt-network-script）。同理，通过定义名为load_config的函数（接收 builder 与 network，返回 builder 配置），也能自定义 TensorRT 的 BuilderConfig。

这一约定的落地在仓库源码中有清晰体现：在 model.py 中，.py扩展名被映射到模型类型trt-network-script；在 loader.py 中，当model_type == "trt-network-script"时，Polygraphy 会通过polygraphy.backend.common的InvokeFromScript加载器来调用脚本中的函数。

约定一：load_network——解析并改造网络的示例脚本

示例目录中的 define_network.py 展示了完整写法：它先解析 ONNX 模型，再在网络上追加一个 Identity 层：

from polygraphy import func from polygraphy.backend.trt import NetworkFromOnnxPath parse_onnx = NetworkFromOnnxPath("identity.onnx") # 定义名为 `load_network` 的函数，polygraphy 会直接用它在命令行中代替模型文件 @func.extend(parse_onnx) def load_network(builder, network, parser): # NOTE: func.extend() 使得该函数的签名变为 `() -> (builder, network, parser)` # 为网络追加一个 identity 层 prev_output = network.get_output(0) network.unmark_output(prev_output) output = network.add_identity(prev_output).get_output(0) network.mark_output(output) # 注意：不需要显式 return 任何东西 —— extend() 已经替我们处理好了

这段代码有两个值得注意的机制要点：

1. func.extend()装饰器：parse_onnx本身是一个返回(builder, network, parser)三元组的"加载器"。func.extend(parse_onnx)会把被装饰函数改写为一个无参、返回(builder, network, parser)的加载器：先执行被扩展的加载器拿到三件套，再进入load_network的函数体对network就地修改。正因为如此，函数体里可以"只修改不返回"。
2. 网络改造方式：先通过network.get_output(0)拿到当前网络第一个输出张量，用unmark_output取消其输出标记，再用add_identity()追加 Identity 层，最后对新张量mark_output。这是"在解析结果上追加层并重定输出"的标准操作范式，适用于任何基于 TensorRT Network API 的图改写。

TIP：如果函数名不叫load_network，可以在模型参数中用冒号显式指定函数名，例如define_network.py:my_func。这一约定同样作用于load_config（见下文）。

约定二：load_config——自定义 Builder 配置的示例脚本

示例目录中的 create_config.py 展示了如何创建一个新的 BuilderConfig 并开启 FP16：

import tensorrt as trt from polygraphy import func from polygraphy.backend.trt import CreateConfig # 定义名为 `load_config` 的函数，polygraphy 用它来创建 builder 配置 @func.extend(CreateConfig()) def load_config(config): # NOTE: func.extend() 使得该函数的签名变为 `(builder, network) -> config` config.set_flag(trt.BuilderFlag.FP16) # 注意：不需要显式 return —— extend() 已经替我们处理好了

与load_network类似，func.extend(CreateConfig())让load_config的签名变为(builder, network) -> config：Polygraphy 先把 builder 与 network 交给内部的CreateConfig()生成一个默认配置对象，再进入函数体让你就地修改，最后自动返回该配置。函数体内通过 TensorRT 官方 APIconfig.set_flag(trt.BuilderFlag.FP16)打开 FP16。

从 config.py 的源码注释可以确认--trt-config-script的完整语义：

指定一个 Python 脚本路径，该脚本定义一个创建 TensorRTIBuilderConfig的函数。函数接收 builder 和 network 作为参数并返回 builder 配置。**当指定该选项时，所有其他 config 参数都会被忽略。**默认情况下 Polygraphy 查找名为load_config的函数；可以通过冒号指定自定义函数名，例如my_custom_script.py:my_func。

这条"优先级最高、全权接管"的语义非常关键：只要传了--trt-config-script，命令行上的--fp16、--int8、--trt-min-shapes等其余配置参数都将被忽略，全部配置逻辑以脚本为准。同理，源码中parse_script_and_func_name（见 config.py）将默认函数名解析为load_config。

另外，load_network与load_config完全可以定义在同一个脚本里，甚至可以从任意脚本或模块中导入这些函数——Polygraphy 的InvokeFromScript机制对来源没有限制，只要最终拿到可调用的函数即可。

TIP：用polygraphy template自动生成脚本模板

与其从零手写脚本，不如先用polygraphy template生成起始模板。仓库的模板实现位于 template/subtool/trt_network.py 与 template/subtool/trt_config.py，CLI 用法如下。

生成网络脚本模板（从零创建网络）：

polygraphy template trt-network -o my_define_network.py

生成的模板核心骨架是：

from polygraphy import func import tensorrt as trt from polygraphy.backend.trt import CreateNetwork @func.extend(CreateNetwork()) def load_network(builder, network): pass # TODO: 在这里搭建网络。该函数不应返回任何内容。

注意：从零创建时函数签名为(builder, network)，对应源码中CreateNetwork()路径（见 trt_network.py）。

如果要从一个已有模型出发、在解析出的 TensorRT 网络上做修改，直接把模型作为参数传给template trt-network即可：

polygraphy template trt-network identity.onnx -o my_define_network.py

此时生成的模板函数签名变为(builder, network, parser)（因为引入了解析器），与本文示例脚本的形态一致。

生成 config 脚本模板：

polygraphy template trt-config -o my_create_config.py

生成的骨架为(builder, network, config)签名（见 trt_config.py）。

用命令行配置项预填充模板。例如要生成一个开启 FP16 的配置脚本：

polygraphy template trt-config --fp16 -o my_create_config.py

因为trt-config模板工具内部订阅了完整的TrtConfigArgs（见 trt_config.py），你在命令行给它的任何配置参数（--fp16、--int8、--trt-min-shapes等）都会以CreateConfig(...)的调用形式被预填充进模板的func.extend(...)装饰器中，模板生成后只需微调即可。

运行示例：从脚本一条命令构建并运行引擎

在示例目录（即04_defining_a_tensorrt_network_or_config_manually/下，与identity.onnx同目录）执行：

1. 直接运行define_network.py定义的网络：

polygraphy run --trt define_network.py --model-type=trt-network-script

--model-type=trt-network-script是必须的（或依赖.py扩展名的自动推断，见 model.py）。Polygraphy 会调用load_network得到 builder/network/parser，然后完成引擎构建与推理。此时identity.onnx解析出的网络末尾被追加了 Identity 层，最终网络输出是 Identity 层的输出。

2. 在步骤 1 的基础上，套用create_config.py定义的 FP16 BuilderConfig：

polygraphy run --trt define_network.py --model-type=trt-network-script --trt-config-script=create_config.py

此时引擎会以 FP16 模式构建。正如前文所述，--trt-config-script会接管全部配置逻辑。

TIP：同样的方式适用于polygraphy convert——如果想只构建引擎而不运行推理：

polygraphy convert --trt define_network.py --model-type=trt-network-script --trt-config-script=create_config.py -o my_engine.engine

即load_network/load_config约定对所有接受模型参数的 Polygraphy CLI 工具通用。

底层机制：Polygraphy 如何加载脚本函数

结合仓库源码，可以还原这条调用链：

1. 模型类型判定：ModelArgs在解析模型参数时，把.py扩展名映射为trt-network-script（见 model.py），并将冒号后的部分解析为函数名extra_model_info，默认值为load_network（见 model.py）。同时，trt-network-script类型要求脚本必须以.py结尾且文件存在，否则直接报 critical 错误。
2. 网络加载：在TrtLoadNetworkArgs.add_to_script_impl中，当模型类型为trt-network-script时，生成InvokeFromScript(model_file, name=network_func_name)加载器（见 loader.py）。InvokeFromScript负责动态导入脚本、按名称取函数并调用，这正是"函数名可以来自任意脚本/模块"的实现基础。
3. 配置加载：在TrtConfigArgs.add_to_script_impl中，若指定了--trt-config-script，同样用InvokeFromScript包装脚本与函数名（见 config.py）；未指定时才走默认的CreateConfig路径并把命令行参数逐一映射进去。
4. 装饰器契约：func.extend保证了用户函数与加载器之间的参数/返回值契约——网络侧是() -> (builder, network, parser)，配置侧是(builder, network) -> config，用户函数体因此可以"只改不返"。

使用注意事项小结

• 脚本必须是.py文件，且load_network默认无参；从零创建网络时函数签名为(builder, network)，解析模型时为(builder, network, parser)。
• 函数名可自定义：模型参数或--trt-config-script值后用冒号指定，如define_network.py:my_func、my_config.py:my_func。
• --trt-config-script会忽略其他所有 config 参数，配置逻辑完全由脚本决定；如需"在 Polygraphy 默认配置基础上微调"，可关注--trt-config-postprocess-script（--trt-cpps，默认函数名postprocess_config，见 config.py），它与--trt-config-script的区别在于前者传入的是已由命令行参数配置好的 config。
• 网络后处理：若不想单独维护load_network脚本，也可以使用实验性的--trt-network-postprocess-script（--trt-npps）对解析出的网络做后处理（默认函数名postprocess）。
• 旧选项--trt-network-func-name与--trt-config-func-name已标记为废弃（自 Polygraphy 0.50.0 起，统一改用冒号语法，见 loader.py 与 config.py）。

延伸阅读

• 示例配套文件：define_network.py、create_config.py、identity.onnx
• 模板工具实现：trt_network.py、trt_config.py
• 参数与加载源码：model.py、loader.py、config.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