使用 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-typetrt-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() 已经替我们处理好了这段代码有两个值得注意的机制要点func.extend()装饰器parse_onnx本身是一个返回(builder, network, parser)三元组的加载器。func.extend(parse_onnx)会把被装饰函数改写为一个无参、返回(builder, network, parser)的加载器先执行被扩展的加载器拿到三件套再进入load_network的函数体对network就地修改。正因为如此函数体里可以只修改不返回。网络改造方式先通过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 并开启 FP16import 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) - configPolygraphy 先把 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.pyCLI 用法如下。生成网络脚本模板从零创建网络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-typetrt-network-script--model-typetrt-network-script是必须的或依赖.py扩展名的自动推断见 model.py。Polygraphy 会调用load_network得到 builder/network/parser然后完成引擎构建与推理。此时identity.onnx解析出的网络末尾被追加了 Identity 层最终网络输出是 Identity 层的输出。2. 在步骤 1 的基础上套用create_config.py定义的 FP16 BuilderConfigpolygraphy run --trt define_network.py --model-typetrt-network-script --trt-config-scriptcreate_config.py此时引擎会以 FP16 模式构建。正如前文所述--trt-config-script会接管全部配置逻辑。TIP同样的方式适用于polygraphy convert——如果想只构建引擎而不运行推理polygraphy convert --trt define_network.py --model-typetrt-network-script --trt-config-scriptcreate_config.py -o my_engine.engine即load_network/load_config约定对所有接受模型参数的 Polygraphy CLI 工具通用。底层机制Polygraphy 如何加载脚本函数结合仓库源码可以还原这条调用链模型类型判定ModelArgs在解析模型参数时把.py扩展名映射为trt-network-script见 model.py并将冒号后的部分解析为函数名extra_model_info默认值为load_network见 model.py。同时trt-network-script类型要求脚本必须以.py结尾且文件存在否则直接报 critical 错误。网络加载在TrtLoadNetworkArgs.add_to_script_impl中当模型类型为trt-network-script时生成InvokeFromScript(model_file, namenetwork_func_name)加载器见 loader.py。InvokeFromScript负责动态导入脚本、按名称取函数并调用这正是函数名可以来自任意脚本/模块的实现基础。配置加载在TrtConfigArgs.add_to_script_impl中若指定了--trt-config-script同样用InvokeFromScript包装脚本与函数名见 config.py未指定时才走默认的CreateConfig路径并把命令行参数逐一映射进去。装饰器契约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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考