torch.fx ¶

注意

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

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

delete_all_unused_submodules()[source]¶

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

如果以下任何一项为真，则模块被认为是“已使用”：1. 它具有已使用的子项 2. 其前向被通过 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)[source]¶: 返回为当前 GraphModule 及其子 GraphModule 生成的 Python 代码

警告

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

recompile()[source]¶

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

注意

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

返回类型: Python代码

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

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

from <folder> import <module_name> 导入

参数

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

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

警告

此 API 为实验性 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)

将生成以下图形

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

构造一个空的图形。

注意

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

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

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

参数

the_function (Callable[..., Any]) – 要调用的函数。可以是任何 PyTorch 运算符、Python 函数或 builtins 或 operator 命名空间的成员。
args (Optional[Tuple[Argument, ...]]) – 要传递给被调用函数的位置参数。
kwargs (Optional[Dict[str, Argument]]) – 要传递给被调用函数的关键字参数
type_expr (Optional[Any]) – 一个可选的类型注释，表示此节点输出将具有的 Python 类型。

返回值

新创建并插入的 call_function 节点。

返回类型

注意

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

注意

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

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

在 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)[source]¶

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

参数

module_name (str) – 要调用的 Module 在 Module 层次结构中的限定名称。例如，如果被跟踪的 Module 具有名为 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)[source]¶

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

参数

op (str) – 该节点的操作码。可以是 ‘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)[source]¶

根据每个节点的用户数量以及节点是否具有任何副作用，从图中删除所有死代码。在调用之前，图必须是拓扑排序的。

参数

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

返回值

图是否因传递而发生更改。

返回类型

示例

在消除死代码之前，下面 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)[source]¶

从 Graph 中删除一个 Node。如果 Graph 中仍然存在该节点的用户，则会抛出异常。

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

注意

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

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

允许快速查询节点。

参数

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

返回值

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

警告

此 API 为实验性 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) – 要从中复制节点的源图。
val_map (Dict[Node, Node]) – 一个字典，其中将填充来自 g 中的节点到 self 中的节点的映射。请注意，可以将 val_map 与其中已经存在的值一起传递，以覆盖某些值的复制。

返回值

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

返回类型

Optional[Union[Tuple[Any, …], List[Any], Dict[str, Any], 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

参数

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

参数

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

整个图的开头插入。

返回值: 将在 __exit__ 上恢复插入点的资源管理器。

注意

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

lint()[source]¶: 对该图运行各种检查，以确保其格式正确。具体包括： - 检查节点是否具有正确的归属权（由该图拥有） - 检查节点是否按拓扑顺序出现 - 如果该图具有拥有 GraphModule，则检查目标是否存在于该 GraphModule 中

注意

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

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

将一个图中的节点复制到另一个图中。 arg_transform 需要将节点所在图的参数转换为 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¶

获取构成此图的节点列表。

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

返回值: 节点的双向链表。请注意，可以对该列表调用 reversed 以切换迭代顺序。

on_generate_code(make_transformer)[source]¶

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

参数

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 为实验性 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 的向后兼容性得到保证。

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

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

参数

name (str) – 输入值的名称。这对应于此 Graph 表示的函数的位置参数的名称。
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]¶: 处理参数，以便可以将其传递到 FX 图中。

警告

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

process_outputs(out)[source]¶: 警告

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

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

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

参数: root_module (str) – 要在其上查找限定名称目标的根模块的名称。这通常是“self”。
返回值: src：表示对象的 Python 源代码 globals：src 中全局名称的字典 -> 它们引用的对象。
返回类型: 一个 PythonCode 对象，包含两个字段

注意

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

set_codegen(codegen)[source]¶

警告

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

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

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

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

返回所有作为此节点输入的节点。这等效于遍历 args 和 kwargs，并仅收集是节点的值。

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

append(x)[source]¶

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

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

注意

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

property args: Tuple[Optional[Union[Tuple[Any, ...], List[Any], Dict[str, Any], slice, range, Node, str, int, float, bool, complex, dtype, Tensor, device, memory_format, layout, OpOverload, SymInt, SymBool, SymFloat]], ...]¶

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

允许对该属性进行赋值。所有使用和用户的核算会在赋值时自动更新。

format_node(placeholder_names=None, maybe_return_typename=None)[source]¶

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

此方法可以用作调试工具，无需任何参数。

此函数还在 Graph 的 __str__ 方法中内部使用。 placeholder_names 和 maybe_return_typename 中的字符串共同构成了此 Graph 所在 GraphModule 中自动生成的 forward 函数的签名。 placeholder_names 和 maybe_return_typename 不应以其他方式使用。

参数

placeholder_names (Optional[List[str]]) – 一个列表，用于存储表示在生成的 forward 函数中的占位符的格式化字符串。仅供内部使用。
maybe_return_typename (Optional[List[str]]) – 一个单元素列表，用于存储表示生成的 forward 函数输出的格式化字符串。仅供内部使用。

返回值

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

返回类型

str

注意

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

insert_arg(idx, arg)[source]¶

在给定索引处插入一个位置参数到参数列表中。

参数

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

注意

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

is_impure()[source]¶

返回此操作是否是不纯的，即其操作是否为占位符或输出，或者是否是调用不纯的 call_function 或 call_module。

返回值

操作是否是不纯的。

返回类型

警告

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

property kwargs: Dict[str, Optional[Union[Tuple[Any, ...], List[Any], Dict[str, Any], slice, range, Node, str, int, float, bool, complex, dtype, Tensor, device, memory_format, layout, OpOverload, SymInt, SymBool, SymFloat]]¶

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

允许对该属性进行赋值。所有使用和用户的核算会在赋值时自动更新。

property next: Node¶

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

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

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

返回规范化后的 Python 目标参数。这意味着 args/kwargs 将与模块/函数的签名相匹配，如果 normalize_to_only_use_kwargs 为真，则以位置顺序返回仅包含 kwargs 的参数。还会填充默认值。不支持仅限位置的参数或可变参数。

支持模块调用。

可能需要 arg_types 和 kwarg_types 来消除重载歧义。

参数

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

返回值

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

返回类型

Optional[ArgsKwargsPair]

警告

此 API 为实验性 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]¶

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

参数

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

返回值

已对此更改进行操作的节点列表。

返回类型

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。调用后，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 函数中的操作。

此方法可以被覆盖以（例如）创建嵌套的已跟踪 GraphModule，或者您在跨越 Module 边界进行跟踪时所需的任何其他行为。

参数

m (Module) – 发出调用的模块
forward (Callable) – 要调用的 Module 的 forward() 方法
args (Tuple) – 模块调用处的参数
kwargs (Dict) – 模块调用处的关键字参数

返回值

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

返回类型

注意

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

create_arg(a)[source]¶

一种方法，用于指定在准备用作 Graph 中节点的参数的值时跟踪的行为。

默认情况下，行为包括

遍历集合类型（例如元组、列表、字典）并在元素上递归调用 create_args。
给定一个 Proxy 对象，返回对底层 IR Node 的引用
给定一个非 Proxy 张量对象，为各种情况发出 IR
- 对于参数，发出一个引用该参数的 get_attr 节点
- 对于非参数张量，将张量存储在一个引用该属性的特殊属性中。

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

参数: a (Any) – 要作为 Graph 中的 Argument 发出的值。
返回值: 将值 a 转换为相应的 Argument
返回类型: Argument

注意

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

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

警告

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

create_node(kind, target, args, kwargs, name=None, type_expr=None)¶

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

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

注意

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

返回类型: 节点

create_proxy(kind, target, args, kwargs, name=None, type_expr=None, proxy_factory_fn=None)¶

根据给定的参数创建节点，然后返回用 Proxy 对象包装的节点。

如果 kind = 'placeholder'，则我们正在创建一个表示函数参数的节点。如果我们需要编码默认参数，则使用 args 元组。对于 placeholder 节点，否则 args 为空。

注意

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

get_fresh_qualname(prefix)[source]¶

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

注意

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

返回类型: str

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

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

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

此方法可以被覆盖以（例如）在查询参数时不返回代理。

参数

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

返回值

来自 getattr 调用的返回值。

警告

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

is_leaf_module(m, module_qualified_name)[source]¶

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

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

参数

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

返回类型

注意

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

iter(obj)¶

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

注意

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

返回类型: 迭代器

keys(obj)¶

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

注意

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

返回类型: 任意

path_of_module(mod)[source]¶

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

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

注意

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

proxy(node)¶

注意

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

返回类型: 代理

to_bool(obj)¶

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

注意

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

返回类型: bool

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

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

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

参数

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

返回值

表示传入 root 语义的 Graph。

返回类型

图

注意

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

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

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

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

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

有两种主要方法可以解决这个问题：1. 将不可跟踪的逻辑分解到顶级函数中，并对其使用 fx.wrap。2. 如果控制流是静态的（即循环次数基于某些超参数），则可以保留代码的原始位置，并将其重构为类似以下内容

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

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

注意

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

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

解释器逐节点执行 FX 图。此模式可用于许多事情，包括编写代码转换以及分析过程。

可以重写解释器类中的方法以自定义执行的行为。根据调用层次结构的可重写方法映射

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

示例

假设我们希望将所有 torch.neg 实例与 torch.sigmoid 实例互换（包括它们的 Tensor 方法等效项）。我们可以像这样子类化解释器

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(n)

    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(n)

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) – 此调用的关键字参数字典

返回类型

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

注意

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

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

执行 call_method 节点并返回结果。

参数

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

返回类型

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

注意

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

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

执行 call_module 节点并返回结果。

参数

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

返回类型