torch.fx #

注意

保证此 API 的向后兼容性。

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

delete_all_unused_submodules()[source]#

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

如果满足以下任一条件，则认为模块是“已使用”的：1. 它有被使用的子节点 2. 它的 forward 通过 call_module 节点直接被调用 3. 它有一个非模块属性通过 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, additional_meta=None)[source]#

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

参数:: additional_meta (list[str] | None) – 要包含在输出中的元键的可选列表。对于列表中的每个键，如果它存在于 node.meta 中，则其值将以“键：值”的格式显示。例如：print_readable(additional_meta=[“seq_nr”])。

警告

此 API 为实验性质，且不保证向后兼容。

recompile()[source]#

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

注意

保证此 API 的向后兼容性。

返回类型:: PythonCode

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

将模块转储到具有 module_name 的 folder 中，以便它可以

使用 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 列表加在一起构成一个有效的 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)

将生成以下 Graph

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]#

在 Graph 中插入一个 call_function Node。 call_function 节点表示对 Python 可调用对象的调用，由 the_function 指定。

参数:

the_function (Callable[..., Any]) – 要调用的函数。可以是任何 PyTorch 运算符、Python 函数，或 builtins 或 operator 命名空间的成员。
args (Optional[Tuple[Argument, ...]]) – 要传递给被调用函数的正向参数。
kwargs (Optional[Dict[str, Argument]]) – 要传递给被调用函数的关键字参数
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)[源码]#

在 Graph 中插入一个 call_method Node（节点）。一个 call_method 节点代表对 args 的第 0 个元素调用给定的方法。

参数:

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

返回:

新创建并插入的 call_method 节点。

返回类型:

注意

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

注意

保证此 API 的向后兼容性。

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

在 Graph 中插入一个 call_module Node。一个 call_module 节点代表对 Module 层级结构中某个 Module 的 forward() 函数的调用。

参数:

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

返回:

新创建并插入的 call_module 节点。

返回类型:

注意

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

注意

保证此 API 的向后兼容性。

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

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

参数:

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

返回:

新创建并插入的节点。

返回类型:

注意

保证此 API 的向后兼容性。

eliminate_dead_code(is_impure_node=None)[源码]#

根据每个节点的使用者数量以及节点是否有副作用，从图中移除所有死代码。在调用此方法之前，图必须已经过拓扑排序。

参数:

is_impure_node (Optional[Callable[[Node], bool]]) – 一个返回以下内容的函数：
None (节点是否为非纯节点。如果此参数为) –
to (，则默认行为是) –
Node.is_impure. (使用) –

返回:

图是否因该 pass 而发生了改变。

返回类型:

示例

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

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

警告

死代码消除包含一些启发式方法，以避免移除有副作用的节点（参见 Node.is_impure），但通常其覆盖范围非常有限。因此，除非你确定 FX 图完全由函数式操作组成，或者你提供了自定义的用于检测副作用节点的函数，否则你不应该假定调用此方法是绝对稳妥的。

注意

保证此 API 的向后兼容性。

erase_node(to_erase)[源码]#

从 Graph 中删除一个 Node。如果该节点在 Graph 中仍有使用者，则抛出异常。

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

注意

保证此 API 的向后兼容性。

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

支持对节点进行快速查询

参数:

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

返回:

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

警告

此 API 为实验性质，且不保证向后兼容。

get_attr(qualified_name, type_expr=None)[源码]#

在 Graph 中插入一个 get_attr 节点。一个 get_attr Node 代表从 Module 层级结构中获取一个属性。

参数:

qualified_name (str) – 要检索的属性的全限定名。例如，如果被追踪的 Module 有一个名为 foo 的子模块，该子模块有一个名为 bar 的子模块，而 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)[源码]#

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

参数:

g (Graph) – 要从中复制节点的源图。
val_map (Dict[Node, Node]) – 一个字典，将被填充为从 g 中的节点到 self 中的节点的映射。请注意，传入 val_map 时可以预先包含一些值，以覆盖某些特定值的复制行为。

返回:

如果 g 具有 output 节点，则返回 self 中现在等效于 g 中输出值的值；否则返回 None。

返回类型:

注意

保证此 API 的向后兼容性。

inserting_after(n=None)[源码]#

设置 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)[源码]#

设置 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()[源码]#: 对该 Graph 运行各种检查，以确保其结构良好。特别是： - 检查 Node 是否拥有正确的所有权（归此图所有） - 检查 Node 是否按拓扑顺序出现 - 如果此 Graph 有所属的 GraphModule，则检查 target 是否存在于该 GraphModule 中

注意

保证此 API 的向后兼容性。

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

将一个节点从一个图复制到另一个图中。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]) – 一个函数，用于将原节点的 args 和 kwargs 中的 Node 参数转换为 self 中等效的参数。在最简单的情况下，它应该从一个将原图节点映射到 self 节点的表中检索值。

返回类型:

注意

保证此 API 的向后兼容性。

property nodes: _node_list#

获取构成此 Graph 的 Node 列表。

请注意，此 Node 列表表示形式是一个双向链表。迭代期间的变动（例如删除节点、添加节点）是安全的。

返回:: 一个双向链表形式的 Node。请注意，可以对此列表调用 reversed 以切换迭代顺序。

on_generate_code(make_transformer)[源码]#

在生成 Python 代码时注册一个转换器（transformer）函数

参数 (Args)

make_transformer (Callable[[Optional[TransformCodeFunc]], TransformCodeFunc])

一个返回要注册的代码转换器的函数。此函数由 on_generate_code 调用以获取代码转换器。

此函数还会接收当前注册的代码转换器作为输入（如果未注册则为 None），以防不希望覆盖它。这对于将代码转换器链式连接在一起非常有用。

返回

一个上下文管理器，当在 with 语句中使用时，会自动恢复以前注册的代码转换器。

示例

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

此函数也可以用作上下文管理器，优点是会自动恢复以前注册的代码转换器

# ... 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)[源码]#

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

参数:

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

注意

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

注意

保证此 API 的向后兼容性。

output_node()[源码]#

警告

此 API 为实验性质，且不保证向后兼容。

返回类型:: Node

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

在 Graph 中插入一个 placeholder 节点。一个 placeholder 代表函数的一个输入。

参数:

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

返回类型:

注意

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

注意

保证此 API 的向后兼容性。

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

注意

保证此 API 的向后兼容性。

process_inputs(*args)[源码]#: 处理参数，以便将它们传递给 FX 图。

警告

此 API 为实验性质，且不保证向后兼容。

process_outputs(out)[源码]#: 警告

此 API 为实验性质，且不保证向后兼容。

python_code(root_module, *, verbose=False, include_stride=False, include_device=False, colored=False, expanded_def=False, record_func=False, additional_meta=None)[源码]#

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

参数:: root_module (str) – 根模块的名称，用于查找限定名目标。通常为 ‘self’。
返回:: src: 代表该对象的 Python 源代码。globals: src 中全局名称与其引用对象之间的映射字典。
返回类型:: 一个 PythonCode 对象，包含两个字段

注意

保证此 API 的向后兼容性。

set_codegen(codegen)[源码]#

警告

此 API 为实验性质，且不保证向后兼容。

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

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

placeholder 代表函数输入。name 属性指定该值将采用的名称。target 同样是参数的名称。args 包含以下内容之一：1) 无内容，或 2) 代表函数输入默认参数的单个参数。kwargs 不受关注。占位符对应于图打印输出中的函数参数（例如 x）。
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 输入的所有 Node。这等同于遍历 args 和 kwargs 并仅收集作为 Node 的值。

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

append(x)[源码]#

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

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

注意

保证此 API 的向后兼容性。

此 Node 的参数元组。对参数的解释取决于节点的 opcode。更多信息请参见 Node 的文档字符串。

允许对此属性进行赋值。赋值时，所有对使用情况（uses）和使用者（users）的记账都会自动更新。

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

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

此方法可在不带参数的情况下作为调试实用程序使用。

此函数还在 Graph 的 __str__ 方法内部使用。通过 placeholder_names 和 maybe_return_typename 中的字符串，共同构成了该 Graph 所属 GraphModule 中自动生成的 forward 函数的签名。不应在其他情况下使用 placeholder_names 和 maybe_return_typename。

参数:

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

返回:

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

返回类型:

str

注意

保证此 API 的向后兼容性。

insert_arg(idx, arg)[源码]#

将位置参数插入到给定索引的参数列表中。

参数:

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

注意

保证此 API 的向后兼容性。

is_impure(impure_random=True)[源码]#

返回此操作是否是非纯的（impure），即它的 op 是否是占位符（placeholder）或输出（output），或者它是否是一个非纯的 call_function 或 call_module。

参数:: impure_random (bool) – 是否将随机操作视为非纯操作。
返回:: 该操作是否为非纯操作。
返回类型:: 布尔值

警告

此 API 为实验性质，且不保证向后兼容。

此 Node 的关键字参数字典。对参数的解释取决于节点的 opcode。更多信息请参见 Node 的文档字符串。

允许对此属性进行赋值。赋值时，所有对使用情况（uses）和使用者（users）的记账都会自动更新。

property next: Node#

返回节点链表中的下一个 Node。

返回:: 节点链表中的下一个 Node。

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

返回 Python 目标的规范化参数。这意味着 args/kwargs 将与模块/函数的签名进行匹配，如果 normalize_to_only_use_kwargs 为真，则按位置顺序排他地返回 kwargs。此方法还会填充默认值。不支持仅限位置的参数（positional-only parameters）或可变长参数（varargs parameters）。

支持模块调用。

可能需要 arg_types 和 kwarg_types 以便消除重载的歧义。

参数:

root (torch.nn.Module) – 用于解析模块目标的模块。
arg_types (Optional[Tuple[Any]]) – 位置参数的类型元组
kwarg_types (Optional[Dict[str, Any]]) – 关键字参数的类型字典
normalize_to_only_use_kwargs (bool) – 是否规范化为仅使用关键字参数。

返回:

返回命名元组 ArgsKwargsPair，如果不成功则返回 None。

返回类型:

ArgsKwargsPair | None

警告

此 API 为实验性质，且不保证向后兼容。

prepend(x)[源码]#

在图的节点列表中将 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=None, *, propagate_meta=False)[源码]#

在图中将所有对 self 的使用替换为节点 replace_with。

参数:

replace_with (Node) – 用于替换所有对 self 使用的节点。
delete_user_cb (Callable) – 被调用的回调函数，用于确定是否应移除 self 节点的给定使用者。
propagate_meta (bool) – 是否将原始节点 .meta 字段上的所有属性复制到替换节点上。为了安全起见，仅当替换节点尚不存在 .meta 字段时，执行此操作才有效。

返回:

进行了此更改的 Node 列表。

返回类型:

list[Node]

注意

保证此 API 的向后兼容性。

replace_input_with(old_input, new_input)[源码]#

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

参数:

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

注意

保证此 API 的向后兼容性。

property stack_trace: str | None#

返回在追踪期间记录的 Python 调用栈轨迹（如果有）。当使用 fx.Tracer 追踪时，此属性通常由 Tracer.create_proxy 填充。要在追踪期间出于调试目的记录栈轨迹，请在 Tracer 实例上设置 record_stack_traces = True。当使用 dynamo 追踪时，此属性默认由 OutputGraph.create_proxy 填充。

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

update_arg(idx, arg)[源码]#

更新现有的位置参数以包含新值 arg。调用后，self.args[idx] == arg。

参数:

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

注意

保证此 API 的向后兼容性。

update_kwarg(key, arg)[源码]#

更新现有的关键字参数以包含新值 arg。调用后，self.kwargs[key] == arg。

参数:

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

注意

保证此 API 的向后兼容性。

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

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

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

注意

保证此 API 的向后兼容性。

call_module(m, forward, args, kwargs)[源码]#

该方法指定了当 Tracer 遇到对 nn.Module 实例的调用时的行为。

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

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

参数:

m (Module) – 正在发射调用节点的模块
forward (Callable) – 要被调用的 Module 的 forward() 方法
args (Tuple) – 模块调用点的位置参数
kwargs (Dict) – 模块调用点的关键字参数

返回:

Module 调用的返回值。如果发出了 call_module 节点，则这是一个 Proxy 值。否则，它是 Module 调用返回的任何值。

返回类型:

注意

保证此 API 的向后兼容性。

create_arg(a)[source]#

一种指定在准备用作 Graph 中节点参数的值时的追踪（tracing）行为的方法。

默认情况下，行为包括：

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

可以重写此方法以支持更多类型。

参数:: a (Any) – 将在 Graph 中作为 Argument 发出的值。
返回:: 转换为适当 Argument 的值 a。
返回类型:: 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]#

根据给定的 target、args、kwargs 和 name 插入一个图节点。

可以重写此方法，以对节点创建中使用的值进行额外的检查、验证或修改。例如，用户可能希望禁止记录就地（in-place）操作。

注意

保证此 API 的向后兼容性。

返回类型:: Node

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

根据给定参数创建一个 Node，然后返回封装在 Proxy 对象中的该 Node。

如果 kind = 'placeholder'，那么我们正在创建一个代表函数参数的 Node。如果我们需要编码默认参数，我们使用 args 元组。对于其他的 placeholder Node，args 为空。

注意

保证此 API 的向后兼容性。

get_fresh_qualname(prefix)[source]#

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

注意

保证此 API 的向后兼容性。

返回类型:: str

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

当我们在调用 nn.Module 实例时调用 getattr，该方法指定此 Tracer 的行为。

默认情况下，其行为是返回该属性的 proxy 值。它还会将 proxy 值存储在 parameter_proxy_cache 中，以便后续调用重用该 proxy，而不是创建新的 proxy。

可以重写此方法，例如在查询参数时不返回 proxy。

参数:

attr (str) – 正在查询的属性名称
attr_val (Any) – 属性的值
parameter_proxy_cache (Dict[str, Any]) – 属性名称到 proxy 的映射缓存

返回:

getattr 调用的返回值。

警告

此 API 为实验性质，且不保证向后兼容。

is_leaf_module(m, module_qualified_name)[source]#

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

叶子模块是出现在 IR 中的原子单元，由 call_module 调用引用。默认情况下，PyTorch 标准库命名空间 (torch.nn) 中的模块都是叶子模块。除非通过此参数另行指定，否则所有其他模块都会被追踪并记录其组成的算子。

参数:

m (Module) – 正在查询的模块
module_qualified_name (str) – 此模块到 root 的路径。例如，如果您有一个模块层次结构，其中子模块 foo 包含子模块 bar，而 bar 又包含子模块 baz，则该模块在此处将以限定名称 foo.bar.baz 出现。

返回类型:

注意

保证此 API 的向后兼容性。

iter(obj)[source]#

当正在对 proxy 对象进行迭代时调用，例如：: 在控制流中使用时。通常我们不知道该怎么做，因为我们不知道 proxy 的值，但自定义 tracer 可以使用 create_node 向图节点附加更多信息，并可以选择返回一个迭代器。

注意

保证此 API 的向后兼容性。

返回类型:: 迭代器 (Iterator)

keys(obj)[source]#

当调用 proxy 对象的 keys() 方法时被调用。: 这是在 proxy 上调用 ** 时发生的情况。如果 ** 应该在您的自定义 tracer 中工作，这应该返回一个迭代器。

注意

保证此 API 的向后兼容性。

返回类型:: 任何

path_of_module(mod)[source]#

在 root 的模块层级结构中查找 mod 限定名称的辅助方法。例如，如果 root 有一个名为 foo 的子模块，而 foo 又有一个名为 bar 的子模块，将 bar 传入此函数将返回字符串 "foo.bar"。

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

注意

保证此 API 的向后兼容性。

proxy(node)[source]#

注意

保证此 API 的向后兼容性。

返回类型:: Proxy

to_bool(obj)[source]#

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

注意

保证此 API 的向后兼容性。

返回类型:: 布尔值

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

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

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

参数:

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

返回:

表示传入的 root 语义的 Graph。

返回类型:

Graph

注意

保证此 API 的向后兼容性。

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

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

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

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

解决此问题主要有两种方法：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)[source]#

Interpreter（解释器）逐节点执行 FX 图。此模式对许多事情都很有用，包括编写代码转换以及分析传递（analysis passes）。

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 is 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)[source]#: 通过解释方式运行 module 并返回结果。这使用了“boxed”调用约定，即您传递一个参数列表，该列表将被解释器清除。这确保了输入张量被及时释放。

注意

保证此 API 的向后兼容性。

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

执行 call_function 节点并返回结果。

参数:

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

返回类型:

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

注意

保证此 API 的向后兼容性。

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

执行 call_method 节点并返回结果。

参数:

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

返回类型:

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

注意

保证此 API 的向后兼容性。

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

执行 call_module 节点并返回结果。

参数:

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

返回类型: