Python PyTorch Proxy用法及代码示例

本文简要介绍python语言中 torch.fx.Proxy 的用法。

用法:

class torch.fx.Proxy(node, tracer=None)

Proxy 对象是 Node 包装器，在符号跟踪期间流经程序，并记录它们接触到不断增长的 FX Graph 的所有操作(torch 函数调用、方法调用、运算符)。

如果您正在执行图形转换，您可以将自己的 Proxy 方法包装在原始 Node 周围，以便您可以使用重载的运算符向 Graph 添加其他内容。

Proxy 对象无法迭代。换句话说，如果在循环中使用 Proxy 或作为 *args /**kwargs 函数参数，符号跟踪器将引发错误。

解决这个问题的主要方法有两种： 1. 将无法追踪的逻辑分解为顶级函数并在其上使用fx.wrap。 2. 如果控制流是静态的(即循环行程计数基于某个超参数)，代码可以保持在其原始位置并重构为：

for i in range(self.some_hyperparameter):
    indexed_item = proxied_value[i]

有关代理内部的更详细说明，请查看 torch/fx/OVERVIEW.md 中的 “Proxy” 部分

注意

保证此 API 的向后兼容性。