torch.fx #

注意

此 API 的向后兼容性已得到保证。

property code: str#: 返回从此 GraphModule 底层的 Graph 生成的 Python 代码。

delete_all_unused_submodules()[source]#

从 self 中删除所有未使用的子模块。

当满足以下任一条件时，模块被认为是“已使用”：1. 它有已使用的子模块 2. 它的 forward 被直接通过 call_module 节点调用 3. 它有一个非 Module 属性，该属性从 get_attr 节点使用。

此方法可用于清理 nn.Module，而无需手动对每个未使用的子模块调用 delete_submodule。

注意

此 API 的向后兼容性已得到保证。

delete_submodule(target)[source]#

从 self 中删除给定的子模块。

如果 target 不是有效目标，则不会删除该模块。

参数

target (str) – 新子模块的完全限定名称（请参阅 nn.Module.get_submodule 中的示例，了解如何指定完全限定名称）。

返回

目标字符串是否引用了: 我们要删除的子模块。返回值为 False 意味着 target 不是对子模块的有效引用。

返回类型

注意

此 API 的向后兼容性已得到保证。

property graph: Graph#: 返回此 GraphModule 底层的 Graph。

print_readable(print_output=True, include_stride=False, include_device=False, colored=False, *, fast_sympy_print=False, expanded_def=False)[source]#

返回当前 GraphModule 及其子 GraphModule 生成的 Python 代码。

警告

此 API 是实验性的，并且不向后兼容。

recompile()[source]#

从其 graph 属性重新编译此 GraphModule。在编辑了包含的 graph 之后应调用此方法，否则此 GraphModule 的生成代码将过时。

注意

此 API 的向后兼容性已得到保证。

返回类型: PythonCode

to_folder(folder, module_name='FxModule')[source]#

将模块转储到 folder 中，使用 module_name，以便可以

使用 from <folder> import <module_name> 导入。

参数 (Args)

folder (Union[str, os.PathLike]): 要将代码写入的文件夹。

module_name (str): 在
写出代码时用于 Module 的顶级名称。

警告

此 API 是实验性的，并且不向后兼容。

class torch.fx.Graph(owning_module=None, tracer_cls=None, tracer_extras=None)[source]#

Graph 是 FX 中间表示使用的主要数据结构。它由一系列 Node 组成，每个 Node 代表一个调用点（或其他语法结构）。Node 的列表共同构成一个有效的 Python 函数。

例如，以下代码：

import torch
import torch.fx


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(3, 4))
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return torch.topk(
            torch.sum(self.linear(x + self.linear.weight).relu(), dim=-1), 3
        )


m = MyModule()
gm = torch.fx.symbolic_trace(m)

将产生以下图：

print(gm.graph)

graph(x):
    %linear_weight : [num_users=1] = self.linear.weight
    %add_1 : [num_users=1] = call_function[target=operator.add](args = (%x, %linear_weight), kwargs = {})
    %linear_1 : [num_users=1] = call_module[target=linear](args = (%add_1,), kwargs = {})
    %relu_1 : [num_users=1] = call_method[target=relu](args = (%linear_1,), kwargs = {})
    %sum_1 : [num_users=1] = call_function[target=torch.sum](args = (%relu_1,), kwargs = {dim: -1})
    %topk_1 : [num_users=1] = call_function[target=torch.topk](args = (%sum_1, 3), kwargs = {})
    return topk_1

有关 Graph 中操作语义的信息，请参阅 Node。

注意

此 API 的向后兼容性已得到保证。

__init__(owning_module=None, tracer_cls=None, tracer_extras=None)[source]#

构造一个空的 Graph。

注意

此 API 的向后兼容性已得到保证。

call_function(the_function, args=None, kwargs=None, type_expr=None, name=None)[source]#

将 call_function Node 插入到 Graph 中。call_function 节点表示对由 the_function 指定的 Python 可调用对象的调用。

参数

the_function (Callable[..., Any]) – 要调用的函数。可以是任何 PyTorch 运算符、Python 函数或 builtins 或 operator 命名空间中的成员。
args (Optional[Tuple[Argument, ...]]) – 要传递给被调用函数的 positional 参数。
kwargs (Optional[Dict[str, Argument]]) – 要传递给被调用函数的 keyword 参数。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。
name (Optional[str]) – 节点的名称。如果未指定，则设置为 None。

返回

新创建并插入的 call_function 节点。

返回类型

注意

此方法的插入点和类型表达式规则与 Graph.create_node() 相同。

注意

此 API 的向后兼容性已得到保证。

call_method(method_name, args=None, kwargs=None, type_expr=None)[source]#

将 call_method Node 插入到 Graph 中。call_method 节点表示对 args 的第 0 个元素上的给定方法进行的调用。

参数

method_name (str) – 要应用于 self 参数的方法名称。例如，如果 args[0] 是表示 Tensor 的 Node，那么要在该 Tensor 上调用 relu()，请将 relu 传递给 method_name。
args (Optional[Tuple[Argument, ...]]) – 要传递给被调用方法的 positional 参数。请注意，这应该包含一个 self 参数。
kwargs (Optional[Dict[str, Argument]]) – 要传递给被调用方法的 keyword 参数。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。

返回

新创建并插入的 call_method 节点。

返回类型

注意

此方法的插入点和类型表达式规则与 Graph.create_node() 相同。

注意

此 API 的向后兼容性已得到保证。

call_module(module_name, args=None, kwargs=None, type_expr=None)[source]#

将一个call_module Node 插入到Graph 中。call_module 节点表示在Module 层次结构中调用Module 的 forward() 函数。

参数

module_name (str) – 要调用的Module 在Module 层次结构中的限定名称。例如，如果跟踪的Module 包含一个名为foo 的子模块，该子模块又包含一个名为bar 的子模块，则应将限定名称foo.bar 作为module_name 传递，以调用该模块。
args (Optional[Tuple[Argument, ...]]) – 要传递给被调用方法的 positional 参数。请注意，这不应包含self 参数。
kwargs (Optional[Dict[str, Argument]]) – 要传递给被调用方法的 keyword 参数。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。

返回

新创建并插入的call_module 节点。

返回类型

注意

此方法的插入点和类型表达式规则与 Graph.create_node() 相同。

注意

此 API 的向后兼容性已得到保证。

create_node(op, target, args=None, kwargs=None, name=None, type_expr=None)[source]#

创建一个Node 并将其插入到当前插入点的Graph 中。请注意，当前插入点可以通过Graph.inserting_before() 和Graph.inserting_after() 来设置。

参数

op (str) – 此 Node 的操作码。可以是 ‘call_function’、‘call_method’、‘get_attr’、‘call_module’、‘placeholder’ 或 ‘output’。这些操作码的语义在Graph 文档字符串中有所描述。
args (Optional[Tuple[Argument, ...]]) – 是此节点的参数元组。
kwargs (Optional[Dict[str, Argument]]) – 此 Node 的 kwargs。
name (Optional[str]) – Node 的可选字符串名称。这会影响在生成的 Python 代码中分配给该值的名称。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。

返回

新创建并插入的节点。

返回类型

注意

此 API 的向后兼容性已得到保证。

eliminate_dead_code(is_impure_node=None)[source]#

从图中删除所有死代码，基于每个节点的 use 数量以及节点是否具有任何 side effects。调用前必须对图进行拓扑排序。

参数

is_impure_node (Optional[Callable[[Node], bool]]) – 一个返回
None (节点是否为 impure。如果是) –
to (则默认行为是) –
Node.is_impure. (使用) –

返回

图是否因该过程而改变。

返回类型

示例

在消除死代码之前，下面 a = x + 1 中的 a 没有 use，因此可以从图中消除而不产生任何影响。

def forward(self, x):
    a = x + 1
    return x + self.attr_1

消除死代码后，a = x + 1 已被删除，其余的 forward 保持不变。

def forward(self, x):
    return x + self.attr_1

警告

死代码消除有一些启发式方法来避免删除具有 side effects 的节点 (参见 Node.is_impure)，但总的来说覆盖率非常差，因此您应该假定调用此方法不是安全的，除非您知道您的 FX 图完全由函数式操作组成，或者您提供了自己的自定义函数来检测具有 side effects 的节点。

注意

此 API 的向后兼容性已得到保证。

erase_node(to_erase)[source]#

从Graph 中删除一个Node。如果该节点在Graph 中仍有 use，则会引发异常。

参数: to_erase (Node) – 要从Graph 中删除的Node。

注意

此 API 的向后兼容性已得到保证。

find_nodes(*, op, target=None, sort=True)[source]#

允许快速查询节点。

参数

op (str) – 操作的名称。
target (Optional[Target]) – 节点的 target。对于 call_function，target 是必需的。对于其他 op，target 是可选的。
sort (bool) – 是否按节点在图上出现的顺序返回节点。

返回

具有请求的操作和 target 的节点的可迭代对象。

警告

此 API 是实验性的，并且不向后兼容。

get_attr(qualified_name, type_expr=None)[source]#

将一个get_attr 节点插入到 Graph 中。get_attr Node 表示从Module 层次结构中获取一个属性。

参数

qualified_name (str) – 要检索的属性的完全限定名称。例如，如果跟踪的 Module 包含一个名为foo 的子模块，该子模块包含一个名为bar 的子模块，该子模块有一个名为baz 的属性，则应将限定名称foo.bar.baz 作为qualified_name 传递。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。

返回

新创建并插入的get_attr 节点。

返回类型

注意

此方法与Graph.create_node 具有相同的插入点和类型表达式规则。

注意

此 API 的向后兼容性已得到保证。

graph_copy(g, val_map, return_output_node=False)[source]#

将给定图中的所有节点复制到self 中。

参数

g (Graph) – 要从中复制 Nodes 的源图。
val_map (Dict[Node, Node]) – 一个字典，将填充从g 中的节点到self 中的节点的映射。请注意，可以传入带有值的val_map 来覆盖某些值的复制。

返回

self 中与g 中的输出值等效的值，如果g 有一个output 节点。否则为None。

返回类型

Optional[Union[tuple[‘Argument’, …], Sequence[Argument], Mapping[str, Argument], slice, range, Node, str, int, float, bool, complex, dtype, Tensor, device, memory_format, layout, OpOverload, SymInt, SymBool, SymFloat]]

注意

此 API 的向后兼容性已得到保证。

inserting_after(n=None)[source]#

设置 create_node 和伴随方法将插入到图中的点。

在 `with` 语句中使用时，这将临时设置插入点，然后在 `with` 语句退出时将其恢复。

with g.inserting_after(n):
    ...  # inserting after node n
...  # insert point restored to what it was previously
g.inserting_after(n)  #  set the insert point permanently

参数 (Args)

n (Optional[Node]): 在其之前插入的节点。如果为 None，则将在图的开头之后插入。
整个图的开头。

返回: 一个资源管理器，它将在 `__exit__` 时恢复插入点。

注意

此 API 的向后兼容性已得到保证。

inserting_before(n=None)[source]#

设置 create_node 和伴随方法将插入到图中的点。

在 `with` 语句中使用时，这将临时设置插入点，然后在 `with` 语句退出时将其恢复。

with g.inserting_before(n):
    ...  # inserting before node n
...  # insert point restored to what it was previously
g.inserting_before(n)  #  set the insert point permanently

参数 (Args)

n (Optional[Node]): 在其之前插入的节点。如果为 None，则将在图的开头之前插入。
整个图的开头。

返回: 一个资源管理器，它将在 `__exit__` 时恢复插入点。

注意

此 API 的向后兼容性已得到保证。

lint()[source]#: 运行此 Graph 的各种检查，以确保其结构良好。特别是： - 检查 Nodes 是否具有正确的归属（属于此图） - 检查 Nodes 是否按拓扑顺序出现 - 如果此 Graph 具有拥有它的 GraphModule，则检查 target 是否在该 GraphModule 中存在。

注意

此 API 的向后兼容性已得到保证。

node_copy(node, arg_transform=<function Graph.<lambda>>)[source]#

将一个节点从一个图复制到另一个图。`arg_transform` 需要将参数从 node 的图转换为 self 的图。示例

# Copying all the nodes in `g` into `new_graph`
g: torch.fx.Graph = ...
new_graph = torch.fx.graph()
value_remap = {}
for node in g.nodes:
    value_remap[node] = new_graph.node_copy(node, lambda n: value_remap[n])

参数

node (Node) – 要复制到self 中的节点。
arg_transform (Callable[[Node], Argument]) – 一个函数，用于将 node 的 `args` 和 `kwargs` 中的 `Node` 参数转换为 `self` 中等效的参数。在最简单的情况下，这应该从一个将原始图中的 Nodes 映射到 `self` 的表中检索一个值。

返回类型

注意

此 API 的向后兼容性已得到保证。

property nodes: list['Node']#

获取构成此 Graph 的 Nodes 列表。

注意，这个 `Node` 列表表示是一个双向链表。在迭代过程中进行修改（例如，删除一个 Node，添加一个 Node）是安全的。

返回: Nodes 的双向链表。注意，可以对此列表调用 `reversed` 来切换迭代顺序。

on_generate_code(make_transformer)[source]#

在生成 Python 代码时注册一个 transformer 函数。

参数 (Args)

make_transformer (Callable[[Optional[TransformCodeFunc]], TransformCodeFunc])
一个返回代码 transformer 的函数，用于注册。此函数由 `on_generate_code` 调用以获取代码 transformer。

此函数也作为输入提供当前注册的代码 transformer（如果未注册，则为 None），以防不希望覆盖它。这对于链接代码 transformer 很有用。

返回
一个上下文管理器，当在 `with` 语句中使用时，会自动恢复先前注册的代码 transformer。

示例
gm: fx.GraphModule = ...


# This is a code transformer we want to register. This code
# transformer prepends a pdb import and trace statement at the very
# beginning of the generated torch.fx code to allow for manual
# debugging with the PDB library.
def insert_pdb(body):
    return ["import pdb; pdb.set_trace()\n", *body]


# Registers `insert_pdb`, and overwrites the current registered
# code transformer (given by `_` to the lambda):
gm.graph.on_generate_code(lambda _: insert_pdb)

# Or alternatively, registers a code transformer which first
# runs `body` through existing registered transformer, then
# through `insert_pdb`:
gm.graph.on_generate_code(
    lambda current_trans: (
        lambda body: insert_pdb(
            current_trans(body) if current_trans else body
        )
    )
)

gm.recompile()
gm(*inputs)  # drops into pdb
此函数还可以用作上下文管理器，其优点是可以自动恢复先前注册的代码 transformer。
# ... continue from previous example

with gm.graph.on_generate_code(lambda _: insert_pdb):
    # do more stuff with `gm`...
    gm.recompile()
    gm(*inputs)  # drops into pdb

# now previous code transformer is restored (but `gm`'s code with pdb
# remains - that means you can run `gm` with pdb here too, until you
# run next `recompile()`).

警告

此 API 是实验性的，并且不向后兼容。

output(result, type_expr=None)[source]#

将一个 `output` Node 插入到Graph 中。`output` 节点表示 Python 代码中的 `return` 语句。`result` 是应返回的值。

参数

result (Argument) – 要返回的值。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出将具有的 Python 类型。

注意

此方法与Graph.create_node 具有相同的插入点和类型表达式规则。

注意

此 API 的向后兼容性已得到保证。

output_node()[source]#

警告

此 API 是实验性的，并且不向后兼容。

返回类型: Node

placeholder(name, type_expr=None, default_value)[source]#

将一个 `placeholder` 节点插入到 Graph 中。`placeholder` 表示函数输入。

参数

name (str) – 输入值的名称。这对应于此 `Graph` 所代表的函数的 positional 参数的名称。
type_expr (Optional[Any]) – 一个可选的类型注解，表示此节点输出的 Python 类型。在某些情况下，这对于正确的代码生成是必需的（例如，当函数随后在 TorchScript 编译中使用时）。
default_value (Any) – 此函数参数应采用的默认值。注意：为了允许将 `None` 作为默认值，应将 `inspect.Signature.empty` 作为此参数传递，以指定该参数*没有*默认值。

返回类型

注意

此方法与Graph.create_node 具有相同的插入点和类型表达式规则。

注意

此 API 的向后兼容性已得到保证。

print_tabular()[source]#: 以表格格式打印图的中间表示。请注意，此 API 需要安装 `tabulate` 模块。

注意

此 API 的向后兼容性已得到保证。

process_inputs(*args)[source]#: 处理 args，以便可以将其传递给 FX 图。

警告

此 API 是实验性的，并且不向后兼容。

process_outputs(out)[source]#: 警告

此 API 是实验性的，并且不向后兼容。

python_code(root_module, *, verbose=False, include_stride=False, include_device=False, colored=False, expanded_def=False)[source]#

将此 `Graph` 转换为有效的 Python 代码。

参数: root_module (str) – 用于查找限定名称 target 的根模块的名称。通常是 ‘self’。
返回: src：表示对象的 Python 源代码 globals：`src` 中的全局名称字典 -> 它们引用的对象。
返回类型: 一个 PythonCode 对象，包含两个字段。

注意

此 API 的向后兼容性已得到保证。

set_codegen(codegen)[source]#

警告

此 API 是实验性的，并且不向后兼容。

class torch.fx.Node(graph, name, op, target, args, kwargs, return_type=None)[source]#

`Node` 是表示 `Graph` 中单个操作的数据结构。在大多数情况下，Nodes 表示对各种实体的调用，例如运算符、方法和 Modules（一些例外包括指定函数输入和输出的节点）。每个 `Node` 都有一个由其 `op` 属性指定的函数。`op` 每个值的 `Node` 语义如下：

`placeholder` 表示函数输入。`name` 属性指定该值将采用的名称。`target` 同样是参数的名称。`args` 持有：1) 空，或 2) 一个表示函数输入的默认参数的单个参数。`kwargs` 是不关心的。
`get_attr` 从模块层次结构中检索参数。`name` 同样是分配给获取结果的名称。`target` 是参数在模块层次结构中的位置的完全限定名称。`args` 和 `kwargs` 是不关心的。
`call_function` 将一个自由函数应用于某些值。`name` 同样是分配给该值的名称。`target` 是要应用的函数。`args` 和 `kwargs` 表示函数的参数，遵循 Python 调用约定。
`call_module` 将模块层次结构中的模块的 `forward()` 方法应用于给定参数。`name` 同上。`target` 是要调用的模块在模块层次结构中的完全限定名称。`args` 和 `kwargs` 表示调用模块的参数，*不包括 self 参数*。
`call_method` 调用一个值上的方法。`name` 同上。`target` 是应用于 `self` 参数的方法的字符串名称。`args` 和 `kwargs` 表示调用模块的参数，*包括 self 参数*。
`output` 在其 `args[0]` 属性中包含跟踪函数的输出。这对应于 Graph 打印输出中的“return”语句。

注意

此 API 的向后兼容性已得到保证。

property all_input_nodes: list['Node']#

返回作为此 Node 输入的所有 Nodes。这相当于迭代 `args` 和 `kwargs` 并仅收集是 Nodes 的值。

返回: 出现在此 `Node` 的 `args` 和 `kwargs` 中的 `Nodes` 列表，按该顺序。

append(x)[source]#

在图节点列表中在此节点之后插入 `x`。等同于 `self.next.prepend(x)`。

参数: x (Node) – 要放在此节点之后的节点。必须是同一图的成员。

注意

此 API 的向后兼容性已得到保证。

property args: tuple[Union[tuple['Argument', ...], collections.abc.Sequence['Argument'], collections.abc.Mapping[str, 'Argument'], slice, range, torch.fx.node.Node, str, int, float, bool, complex, torch.dtype, torch.Tensor, torch.device, torch.memory_format, torch.layout, torch._ops.OpOverload, torch.SymInt, torch.SymBool, torch.SymFloat, NoneType], ...]#

此 `Node` 的参数元组。参数的解释取决于节点的 op。有关更多信息，请参阅 Node 文档字符串。

允许对此属性进行赋值。所有 use 和 users 的计数都会在赋值时自动更新。

format_node(placeholder_names=None, maybe_return_typename=None, *, include_tensor_metadata=False)[source]#

返回 `self` 的描述性字符串表示。

此方法可作为调试实用程序使用，无需参数。

此函数还在 `Graph` 的 `__str__` 方法中用作内部助手。`placeholder_names` 和 `maybe_return_typename` 中的字符串共同构成了此 Graph 外部 GraphModule 中自动生成的 `forward` 函数的签名。`placeholder_names` 和 `maybe_return_typename` 不得在其他地方使用。

参数

placeholder_names (Optional[list[str]]) – 一个列表，将存储表示生成的 `forward` 函数中 placeholder 的格式化字符串。仅内部使用。
maybe_return_typename (Optional[list[str]]) – 一个单元素列表，将存储表示生成的 `forward` 函数输出的格式化字符串。仅内部使用。
include_tensor_metadata (bool) – 是否包含张量元数据。

返回

如果 1) 我们正在使用 `format_node` 作为内部助手: 在 `Graph` 的 `__str__` 方法中，并且 2) `self` 是一个 placeholder Node，则返回 `None`。否则，返回当前 Node 的描述性字符串表示。

返回类型

str

注意

此 API 的向后兼容性已得到保证。

insert_arg(idx, arg)[source]#

将一个 positional 参数插入到具有给定索引的参数列表中。

参数

idx (int) – 要插入之前的 `self.args` 中的元素的索引。
arg (Argument) – 要插入 `args` 中的新参数值。

注意

此 API 的向后兼容性已得到保证。

is_impure(impure_random=True)[source]#

返回此 op 是否为 impure，即如果其 op 是 placeholder 或 output，或者是一个 impure 的 call_function 或 call_module。

参数: impure_random (bool) – 是否将 rand op 视为 impure。
返回: op 是 impure 还是不是。
返回类型: 布尔值

警告

此 API 是实验性的，并且不向后兼容。

property kwargs: dict[str, Union[tuple['Argument', ...], collections.abc.Sequence['Argument'], collections.abc.Mapping[str, 'Argument'], slice, range, torch.fx.node.Node, str, int, float, bool, complex, torch.dtype, torch.Tensor, torch.device, torch.memory_format, torch.layout, torch._ops.OpOverload, torch.SymInt, torch.SymBool, torch.SymFloat, NoneType]]#

此 `Node` 的关键字参数字典。参数的解释取决于节点的 op。有关更多信息，请参阅 Node 文档字符串。

允许对此属性进行赋值。所有 use 和 users 的计数都会在赋值时自动更新。

property next: Node#

返回链表中下一个 `Node`。

返回: 链表中下一个 `Node`。

normalized_arguments(root, arg_types=None, kwarg_types=None, normalize_to_only_use_kwargs=False)[source]#

返回 Python target 的标准化参数。这意味着 `args/kwargs` 将与 module/functional 的签名匹配，并且如果 `normalize_to_only_use_kwargs` 为 true，则将仅按位置返回 kwargs。还将填充默认值。不支持仅 positional 参数或 varargs 参数。

支持模块调用。

可能需要 `arg_types` 和 `kwarg_types` 来区分重载。

参数

root (torch.nn.Module) – 用于解析模块 target 的模块。
arg_types (Optional[Tuple[Any]]) – arg 的 arg 类型元组。
kwarg_types (Optional[Dict[str, Any]]) – kwargs 的 arg 类型字典。
normalize_to_only_use_kwargs (bool) – 是否规范化为仅使用 kwargs。

返回

返回 NamedTuple ArgsKwargsPair，如果未成功则返回 `None`。

返回类型

Optional[ArgsKwargsPair]

警告

此 API 是实验性的，并且不向后兼容。

prepend(x)[source]#

在图节点列表中在此节点之前插入 x。示例

Before: p -> self
        bx -> x -> ax
After:  p -> x -> self
        bx -> ax

参数: x (Node) – 要放在此节点之前的节点。必须是同一图的成员。

注意

此 API 的向后兼容性已得到保证。

property prev: Node#

返回链表中上一个 `Node`。

返回: 链表中上一个 `Node`。

replace_all_uses_with(replace_with, delete_user_cb=<function Node.<lambda>>, *, propagate_meta=False)[source]#

将 Graph 中 `self` 的所有 use 替换为 Node `replace_with`。

参数

replace_with (Node) – 用于替换 `self` 的所有 use 的节点。
delete_user_cb (Callable) – 用于确定是否应删除 self 节点的给定 use 的回调函数。
propagate_meta (bool) – 是否将原始节点的 .meta 字段中的所有属性复制到替换节点。为安全起见，仅当替换节点尚不具有 .meta 字段时才有效。

返回

在此更改被执行的 Nodes 列表。

返回类型

list[‘Node’]

注意

此 API 的向后兼容性已得到保证。

replace_input_with(old_input, new_input)[source]#

遍历 `self` 的输入节点，并将 `old_input` 的所有实例替换为 `new_input`。

参数

old_input (Node) – 要被替换的旧输入节点。
new_input (Node) – 用于替换 `old_input` 的新输入节点。

注意

此 API 的向后兼容性已得到保证。

property stack_trace: Optional[str]#

返回跟踪期间记录的 Python 堆栈跟踪（如果有）。当使用 fx.Tracer 进行跟踪时，此属性通常由 `Tracer.create_proxy` 填充。要在跟踪期间记录堆栈跟踪以进行调试，请在 `Tracer` 实例上设置 `record_stack_traces = True`。当使用 dynamo 进行跟踪时，此属性将由 `OutputGraph.create_proxy` 默认填充。

stack_trace 的字符串末尾将是最内层的帧。

update_arg(idx, arg)[source]#

使用新值 `arg` 更新现有 positional 参数。调用后，`self.args[idx] == arg`。

参数

idx (int) – 要更新的元素在 `self.args` 中的索引。
arg (Argument) – 要写入 `args` 中的新参数值。

注意

此 API 的向后兼容性已得到保证。

update_kwarg(key, arg)[source]#

使用新值 `arg` 更新现有关键字参数。调用后，`self.kwargs[key] == arg`。

参数

key (str) – 要更新的元素在 `self.kwargs` 中的键。
arg (Argument) – 要写入 `kwargs` 中的新参数值。

注意

此 API 的向后兼容性已得到保证。

class torch.fx.Tracer(autowrap_modules=(math,), autowrap_functions=())[source]#

`Tracer` 是实现 `torch.fx.symbolic_trace` 的符号跟踪功能的类。调用 `symbolic_trace(m)` 等同于 `Tracer().trace(m)`。

Tracer 可以被子类化以覆盖跟踪过程的各种行为。可以覆盖的不同行为在此类方法的文档字符串中进行了描述。

注意

此 API 的向后兼容性已得到保证。

call_module(m, forward, args, kwargs)[source]#

指定此 `Tracer` 在遇到对 `nn.Module` 实例的调用时行为的方法。

默认情况下，行为是检查被调用的模块是否为叶子模块（通过 `is_leaf_module`）。如果是，则在 `Graph` 中发出一个指向 `m` 的 `call_module` 节点。否则，正常调用 `Module`，跟踪其 `forward` 函数中的操作。

可以覆盖此方法来实现，例如，创建嵌套的跟踪 GraphModules，或在跟踪跨越 `Module` 边界时执行任何其他所需行为。

参数

m (Module) – 正在发出调用的模块。
forward (Callable) – 要调用的 `Module` 的 forward() 方法。
args (Tuple) – 模块调用处的 args。
kwargs (Dict) – 模块调用处的 kwargs。

返回

来自 Module 调用的返回值。在发出 `call_module` 节点的情况下，这是一个 `Proxy` 值。否则，它是从 `Module` 调用返回的任何值。

返回类型

注意

此 API 的向后兼容性已得到保证。

create_arg(a)[source]#

用于指定此 `Tracer` 在准备作为 `Graph` 中节点参数的值时的行为的方法。

默认行为包括：

迭代集合类型（例如，tuple、list、dict），并递归调用 `create_args` 来处理元素。
给定 Proxy 对象，返回对底层 IR `Node` 的引用。
给定非 Proxy Tensor 对象，为各种情况发出 IR。
- 对于 Parameter，发出指向该 Parameter 的 `get_attr` 节点。
- 对于非 Parameter Tensor，将该 Tensor 存储在一个特殊的属性中，该属性指向该属性。

可以覆盖此方法以支持更多类型。

参数: a (Any) – 将作为 `Argument` 发出到 `Graph` 中的值。
返回: `a` 转换为适当 `Argument` 的值。
返回类型: Argument

注意

此 API 的向后兼容性已得到保证。

create_args_for_root(root_fn, is_module, concrete_args=None)[source]#: 创建与根模块（`root` Module）的签名对应的 `placeholder` 节点。此方法内省 root 的签名并相应地发出这些节点，还支持 `*args` 和 `**kwargs`。

警告

此 API 是实验性的，并且不向后兼容。

create_node(kind, target, args, kwargs, name=None, type_expr=None)[source]#

根据目标、参数、关键字参数和名称插入图节点。

可以覆盖此方法以执行额外的检查、验证或修改用于节点创建的值。例如，有人可能希望不允许记录就地操作。

注意

此 API 的向后兼容性已得到保证。

返回类型: Node

create_proxy(kind, target, args, kwargs, name=None, type_expr=None, proxy_factory_fn=None)[source]#

根据给定的参数创建 Node，然后将 Node 包装在 Proxy 对象中返回。

如果 kind = ‘placeholder’，那么我们正在创建一个表示函数参数的 Node。如果需要编码默认参数，我们使用 `args` 元组。对于 `placeholder` Nodes，`args` 否则为空。

注意

此 API 的向后兼容性已得到保证。

get_fresh_qualname(prefix)[source]#

为前缀获取一个新名称并返回它。此函数确保它不会与图上的现有属性冲突。

注意

此 API 的向后兼容性已得到保证。

返回类型: str

getattr(attr, attr_val, parameter_proxy_cache)[source]#

指定此 `Tracer` 在调用 `nn.Module` 实例上的 getattr 时行为的方法。

默认情况下，行为是返回属性的代理值。它还将代理值存储在 `parameter_proxy_cache` 中，以便将来的调用将重用该代理而不是创建新代理。

可以覆盖此方法来实现，例如，在查询参数时不返回代理。

参数

attr (str) – 查询的属性名称。
attr_val (Any) – 属性的值。
parameter_proxy_cache (Dict[str, Any]) – 属性名到代理的缓存。

返回

来自 getattr 调用的返回值。

警告

此 API 是实验性的，并且不向后兼容。

is_leaf_module(m, module_qualified_name)[source]#

指定给定 `nn.Module` 是否为“叶子”模块的方法。

叶子模块是 IR 中出现的原子单元，由 `call_module` 调用引用。默认情况下，PyTorch 标准库命名空间（torch.nn）中的模块是叶子模块。所有其他模块都将被跟踪并记录其组成的 op，除非通过此参数另有指定。

参数

m (Module) – 查询的模块。
module_qualified_name (str) – 到此模块根的路径。例如，如果模块层次结构中的子模块 `foo` 包含子模块 `bar`，该子模块又包含子模块 `baz`，则该模块将在此处显示为限定名称 `foo.bar.baz`。

返回类型

注意

此 API 的向后兼容性已得到保证。

iter(obj)[source]#

在迭代代理对象时调用，例如: 在控制流中使用时。通常我们不知道该怎么做，因为我们不知道代理的值，但是自定义跟踪器可以通过 create_node 将更多信息附加到图节点，并可以选择返回一个迭代器。

注意

此 API 的向后兼容性已得到保证。

返回类型: 迭代器。

keys(obj)[source]#

在代理对象调用 keys() 方法时调用。: 这就是当对代理对象执行 ** 时发生的情况。这应该返回一个迭代器，如果 ** 在您的自定义跟踪器中工作的话。

注意

此 API 的向后兼容性已得到保证。

返回类型: 任何

path_of_module(mod)[源]#

辅助方法，用于查找 root 的 Module 层次结构中 mod 的限定名称。例如，如果 root 有一个名为 foo 的子模块，而 foo 又有一个名为 bar 的子模块，则将 bar 传递给此函数将返回字符串 “foo.bar”。

参数: mod (str) – 要为其检索限定名称的 Module。
返回类型: str

注意

此 API 的向后兼容性已得到保证。

proxy(node)[源]#

注意

此 API 的向后兼容性已得到保证。

返回类型: 代理

to_bool(obj)[源]#

当代理对象被转换为布尔值时调用，例如: 在控制流中使用时。通常我们不知道该怎么做，因为我们不知道代理的值，但是自定义跟踪器可以使用 create_node 向图节点附加更多信息，并可以选择返回值。

注意

此 API 的向后兼容性已得到保证。

返回类型: 布尔值

trace(root, concrete_args=None)[源]#

跟踪 root 并返回相应的 FX Graph 表示。 root 可以是 nn.Module 实例或 Python 可调用对象。

请注意，调用此方法后，self.root 可能与此处传入的 root 不同。例如，当将一个自由函数传递给 trace() 时，我们将创建一个 nn.Module 实例作为根并添加嵌入的常量。

参数

root (Union[Module, Callable]) – 要通过跟踪的 Module 或函数。此参数向后兼容。
concrete_args (Optional[Dict[str, any]]) – 不应被视为代理的具体参数。此参数是实验性的，其向后兼容性不有保证。

返回

一个 Graph，表示传入的 root 的语义。

返回类型

Graph

注意

此 API 的向后兼容性已得到保证。

class torch.fx.Proxy(node, tracer=None)[源]#

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

如果您正在进行图变换，您可以将自己的 Proxy 方法包装在原始 Node 上，以便可以使用重载运算符将其他内容添加到 Graph 中。

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

通常有两种方法可以解决此问题：1. 将无法跟踪的逻辑提取到一个顶级函数中，并对其使用 fx.wrap。 2. 如果控制流是静态的（即循环次数取决于某个超参数），则可以将代码保留在原始位置并重构为类似

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

有关 Proxy 内部机制的更详细说明，请参阅 torch/fx/README.md 中的“Proxy”部分。

注意

此 API 的向后兼容性已得到保证。

class torch.fx.Interpreter(module, garbage_collect_values=True, graph=None)[源]#

Interpreter 按节点逐个执行 FX 图。这种模式可用于许多用途，包括编写代码转换和分析传递。

可以覆盖 Interpreter 类中的方法来定制执行行为。可覆盖方法的调用层次图

run()
    +-- run_node
        +-- placeholder()
        +-- get_attr()
        +-- call_function()
        +-- call_method()
        +-- call_module()
        +-- output()

示例

假设我们要将所有 torch.neg 实例与 torch.sigmoid 相互替换（包括它们的 Tensor 方法等效项）。我们可以像这样继承 Interpreter

class NegSigmSwapInterpreter(Interpreter):
    def call_function(
        self, target: Target, args: Tuple, kwargs: Dict
    ) -> Any:
        if target == torch.sigmoid:
            return torch.neg(*args, **kwargs)
        return super().call_function(target, args, kwargs)

    def call_method(self, target: Target, args: Tuple, kwargs: Dict) -> Any:
        if target == "neg":
            call_self, *args_tail = args
            return call_self.sigmoid(*args_tail, **kwargs)
        return super().call_method(target, args, kwargs)


def fn(x):
    return torch.sigmoid(x).neg()


gm = torch.fx.symbolic_trace(fn)
input = torch.randn(3, 4)
result = NegSigmSwapInterpreter(gm).run(input)
torch.testing.assert_close(result, torch.neg(input).sigmoid())

参数

module (torch.nn.Module) – 要执行的模块
garbage_collect_values (bool) – 是否在模块执行期间删除其最后一个使用的值。这可确保在执行过程中实现最佳内存使用。例如，可以通过查看 Interpreter.env 属性来禁用此项，以检查执行中的所有中间值。
graph (Optional[Graph]) – 如果已传递，解释器将执行此图而不是 module.graph，并使用提供的 module 参数来满足任何状态请求。

注意

此 API 的向后兼容性已得到保证。

boxed_run(args_list)[源]#: 通过解释运行 module 并返回结果。这使用了“boxed”调用约定，您需要传递一个参数列表，该列表将被解释器清除。这确保输入张量能被及时释放。

注意

此 API 的向后兼容性已得到保证。

call_function(target, args, kwargs)[源]#

执行 call_function 节点并返回结果。

参数

target (Target) – 此节点的调用目标。有关语义的详细信息，请参阅 Node
args (Tuple) – 此调用的位置参数元组
kwargs (Dict) – 此调用的关键字参数字典

返回类型

Return: Any: 函数调用返回的值

注意

此 API 的向后兼容性已得到保证。

call_method(target, args, kwargs)[源]#

执行 call_method 节点并返回结果。

参数

target (Target) – 此节点的调用目标。有关语义的详细信息，请参阅 Node
args (Tuple) – 此调用的位置参数元组
kwargs (Dict) – 此调用的关键字参数字典

返回类型

Return: Any: 方法调用返回的值

注意

此 API 的向后兼容性已得到保证。

call_module(target, args, kwargs)[源]#

执行 call_module 节点并返回结果。

参数

target (Target) – 此节点的调用目标。有关语义的详细信息，请参阅 Node
args (Tuple) – 此调用的位置参数元组
kwargs (Dict) – 此调用的关键字参数字典

返回类型