深入解析:func.extend 与 func.constantmethod 的设计与实战)
TensorRT Polygraphy 函数助手polygraphy.func深入解析func.extend 与 func.constantmethod 的设计与实战【免费下载链接】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/TensorRTpolygraphy.func是 PolygraphyNVIDIA TensorRT 开源仓库中用于深度学习模型调试与基准测试的工具包提供的一组函数级工具装饰器核心价值在于让开发者以极简代码扩展既有函数尤其是 Polygraphy 的懒加载器的返回值与行为并在内部正确性检查开启时为常量方法提供防变异保护。读完本文你将掌握func.extend的返回值转发、参数转发、返回值覆盖三大语义及其源码级实现原理并能将其直接应用于扩展 TensorRT 的NetworkFromOnnxPath、CreateConfig等加载器写出可复用的构建脚本。polygraphy.func 模块在 Polygraphy 中的定位在 Polygraphy 的 API 参考体系中func模块属于杂项Miscellaneous分组其文档入口为 docs/func/toc.rst标题为 Function Helpers通过 Sphinx 的automodule:: polygraphy.func.func指令从源码 docstring 自动生成 API 文档该入口被 docs/index.rst 的 toctree 收录。模块的实际实现位于 polygraphy/func/func.py并由 polygraphy/func/init.py 通过from polygraphy.func.func import *对外暴露。整个模块只包含两个公开装饰器装饰器作用func.extend(extend_func)用被装饰函数扩展另一个函数转发其返回值可同时转发其入参func.constantmethod标记常量方法防止方法意外修改实例属性受环境变量门控从源码结构看该模块刻意保持极小的公共 API 面仅两个导出符号属于 Polygraphy 内部基础设施——它被大量加载器、模板工具与测试所依赖但对外只提供最精炼的两个装饰器。func.extend函数扩展装饰器func.extend是polygraphy.func的核心其设计目标是解决 Polygraphy 中一个非常典型的痛点Polygraphy 提供的加载器loader返回的往往是 TensorRT 的builder、network、config等底层对象开发者希望在这些对象返回后立即插入自定义修改逻辑却又不想手写样板化的包装代码。extend让被装饰函数天然获得加载器的返回值并自动把这些返回值再透传给调用者。核心语义一返回值自动转发最直观的用法是被extend(x)装饰的函数y其参数就是x的返回值调用y等价于调用x后把返回值交给y处理并最终把x的返回值交给y的调用者。模块 docstring 给出了完整示例def x(a0, a1, a2): rv0 [a0, a1, a2] rv1 None return rv0, rv1 extend(x) def y(rv0, rv1): rv0.append(-1) # 可以像调用 x 一样调用 y并获得 x 经 y 修改后的返回值 rv0, rv1 y(1, 2, 3) assert rv0 [1, 2, 3, -1] assert rv1 is None从语义上讲extend本质上是如下手写包装的语法糖def y(a0, a1, a2): rv0, rv1 x(a0, a1, a2) rv0.append(-1) # 原 y 的函数体 return rv0, rv1关键规则源自模块 docstring 与 func.py 实现若y不返回任何值或返回None则extend会把x的返回值转发给调用者——因此y与x拥有完全一致的对外接口若y返回非None值则该值会交给调用者x的返回值被丢弃即返回值覆盖语义。核心语义二参数转发——同时看到入参与返回值某些场景下被扩展函数需要同时访问x的输入参数和返回值。extend规定只要y在常规参数即x的返回值之前声明了与x完全一致的参数x的所有实参就会被一并转发给y。docstring 中的例子def x(x_arg): return x_arg 1 def y(x_arg, x_ret): # y 现在既能看见 x 的入参也能看见 x 的返回值 assert x_ret x_arg 1 assert y(5) 6有一个重要的注意点文档明确提示如果x原地修改了其参数那么y看到的将是修改后的参数而非原始值——因为参数是按引用传递的。这一能力正是 CHANGELOG.md 中记录的更新点Updatedfunc.extend()to provide a mechanism to accept the input parameters of the extended function也是extend相比普通装饰器最智能的地方。核心语义三返回值元组自动解包extend会自动解包被扩展函数返回的元组。因此下面两种x的写法对y而言完全等价# 写法 A直接返回多个值 def x(a0, a1, a2): rv0 [a0, a1, a2] rv1 None return rv0, rv1 # 写法 B返回一个元组extend 会自动解包y 仍看到 2 个参数 def x(a0, a1, a2): ret (rv0, rv1) return ret这保证y的参数个数始终与解包后的返回值个数一致调用方无需关心x内部返回的是元组还是多个值。源码实现剖析extend的实现位于 func.py 第 29154 行其核心执行逻辑可以拆解为几步调用被扩展函数extend_func_retval extend_func(*args, **kwargs)并借助模块内辅助函数make_iterablefunc.py仅当返回值为tuple时才包成元组统一返回值形态按参数个数分派通过inspect.signature(func).parameters获取被装饰函数y的参数签名然后依据返回值个数 vs 参数个数的组合分三条路径处理y无参数且x返回单值None直接func()空调用返回值个数恰等于y参数个数func(*extend_func_ret_tuple)——即只看返回值的普通转发返回值个数 实参数目等于y参数个数说明y同时声明了x的入参此时把返回值转换为关键字参数追加在**kwargs之后调用func(*args, **kwargs, **ret_kwargs)从而正确处理位置参数/关键字参数混用的场景参数不匹配时报错若三种情况都不满足则通过G_LOGGER.critical抛出致命日志提示y的参数个数与x的返回值个数不匹配并附带返回值的实际类型元组方便定位问题返回值决策func_retval is not None时返回y的返回值覆盖语义否则返回x的返回值转发语义。从实现细节还可以推断出两条使用限制被装饰函数不能使用*args/**kwargs变长参数文档明确 NOTE否则签名推断无法工作装饰器通过functools.wraps(func)保留原函数元信息保证y的__name__、__doc__等属性不被污染。实战扩展 TensorRT 加载器func.extend在 Polygraphy 中最典型的应用场景是扩展官方提供的懒加载器。以 examples/api/03_interoperating_with_tensorrt/example.py 为例NetworkFromOnnxPath(identity.onnx)返回 TensorRT 的builder、network、parser三个对象开发者可以用func.extend直接拿到它们并调用 TensorRT 原生 APIimport tensorrt as trt from polygraphy import func from polygraphy.backend.trt import ( CreateConfig, EngineFromNetwork, NetworkFromOnnxPath, TrtRunner, ) # NetworkFromOnnxPath 返回 builder、network、parser这三个对象会作为参数传入 func.extend(NetworkFromOnnxPath(identity.onnx)) def load_network(builder, network, parser): network.name MyIdentity # 直接用 TensorRT API 修改网络 print(fNetwork name: {network.name}) # 无需 returnextend() 会自动转发返回值 # CreateConfig 返回 IBuilderConfig同样可以直接用 TensorRT API 设置 func.extend(CreateConfig()) def load_config(config): config.set_flag(trt.BuilderFlag.FP16) def main(): # 注意懒加载模式下传入的是函数本身而非调用结果 build_engine EngineFromNetwork(load_network, configload_config) with TrtRunner(build_engine) as runner: outputs runner.infer({x: np.ones(shape(1, 1, 2, 2), dtypenp.float32)}) assert np.array_equal(outputs[y], np.ones(shape(1, 1, 2, 2), dtypenp.float32))同样的模式也出现在 CLI 侧的模板工具与示例中examples/cli/run/04_defining_a_tensorrt_network_or_config_manually/define_network.pyfunc.extend(parse_onnx)使被装饰函数的签名变为() - (builder, network, parser)examples/cli/run/04_defining_a_tensorrt_network_or_config_manually/create_config.pyfunc.extend(CreateConfig())使签名变为(builder, network) - configexamples/cli/run/08_adding_precision_constraints/constrained_network.py用func.extend(parse_network_from_onnx)为网络追加精度约束polygraphy/tools/template/subtool/trt_network.py 与 trt_config.pyPolygraphy 的template命令在生成脚本时会直接向模板追加func.extend(...)代码片段将用户自定义逻辑挂接到网络/配置加载器上。测试验证语义的完整覆盖test_func.py 中的TestExtend测试类对上述语义做了系统性验证可作为行为契约参考test_override_rvy显式返回 2 时x的返回值被丢弃y() 2test_extend_named_parameters支持关键字参数调用y(arg11, arg00)test_extend_0_args_1_rv/test_extend_0_args_2_rvx返回单值/多值时y分别收到 1 个/2 个参数且y()仍原样返回x的返回值test_extend_1_args_0_rv/test_extend_1_args_1_rv/test_extend_2_args_2_rv入参与返回值的各种组合test_extend_can_modify_rv/test_extend_can_modify_rv_objectsy可以原地修改x返回的列表或自定义对象修改结果对调用者可见test_extend_incorrect_num_args参数个数不匹配时抛出PolygraphyException异常信息形如Function: y accepts 1 parameter(s), but needs to accept 2 parameter(s)test_extend_forward_parameters通过pytest.mark.parametrize覆盖全部关键字参数、全部位置参数、混合传参三种模式验证参数转发路径的正确性。这些测试与 func.py 中的实现一一对应读者在二次开发时可以直接以此为行为基线。func.constantmethod常量方法守卫func.constantmethod是模块中的第二个装饰器用于标记常量方法——即不允许修改self实例属性的方法。它解决的是 Polygraphy 内部代码在重构时的防御性问题某些方法按设计不应改变对象状态通过显式标注并在正确性检查开启时运行时校验可以在开发阶段尽早暴露意外的变异。触发条件与门控机制该装饰器最重要的行为特性是仅在POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS环境变量设置为1时才会真正生效对应源码中的config.INTERNAL_CORRECTNESS_CHECKS开关见 func.py。未设置该变量时装饰器直接原样返回func零性能开销。因此它可以安全地批量标注在生产代码上只在内部正确性检查测试/CI阶段发挥约束作用。docstring 中的示例class Dummy: def __init__(self): self.x 1 func.constantmethod def modify_x(self): self.x 2 d Dummy() d.modify_x() # 在正确性检查开启时这会抛出内部错误实现原理与能力边界启用后装饰器在调用前后各执行一次vars(self)快照对比old_dict copy.copy(vars(self))保存方法调用前的实例属性字典执行被装饰方法并用try/finally保证无论是否抛异常都执行对比若vars(self) ! old_dict则通过G_LOGGER.internal_error报告常量方法被变异并输出旧/新状态以便定位。同时文档与实现都明确指出了其能力边界这只是对实例属性引用字典键的最小防护。如果实例持有可变对象引用例如 numpy 数组成员该方法无法保证数组内部数值不被改变——因为引用本身没有变化。测试用例TestConstantMethod覆盖了修改已有属性test_cannot_modify_attrs与新增属性test_cannot_add_attrs两类违反场景均抛出PolygraphyInternalException。使用注意事项与最佳实践综合文档、源码与测试使用polygraphy.func时有以下几点值得注意extend的返回值规则要牢记y返回None时透传x的返回值返回非None时覆盖。若想在y中既不修改返回值也不覆盖直接写空函数体或不写return即可被扩展函数不得使用*args/**kwargs否则签名推断会失效参数转发时注意引用别名x原地修改参数后y看到的是修改后的值constantmethod依赖环境变量POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS1才生效且只防护实例属性的增删改字典层面不防护可变对象内容的变更结合懒加载器使用在 Polygraphy 中传递func.extend装饰后的函数时应传入函数本身而非调用结果如EngineFromNetwork(load_network, configload_config)与懒加载体系保持一致该模块是内部基础设施但 API 稳定docs/index.rst明确警告未在文档中列出的 API 视为内部实现不承诺弃用策略而extend与constantmethod正是被 docs/func/toc.rst 收录的两个公开 API可放心在脚本中使用。小结polygraphy.func用两个装饰器优雅地解决了 Polygraphy 生态中的两个真实问题func.extend让加载器扩展零样板、接口透明返回值自动转发、入参可见、返回值可覆盖func.constantmethod则在不牺牲性能的前提下环境变量门控为内部正确性保驾护航。理解其语义与实现是深入使用 Polygraphy 懒加载 API、编写自定义 TensorRT 构建脚本乃至二次开发 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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考