torch.fx

概述

FX 是一个开发工具包，用于转换 nn.Module 实例。它包含三个主要组件：符号追踪器、中间表示和 Python 代码生成。以下是这些组件的示例演示：

import torch
# Simple module for demonstration
class MyModule(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(3, 4))
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return self.linear(x + self.param).clamp(min=0.0, max=1.0)

module = MyModule()

from torch.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced : torch.fx.GraphModule = symbolic_trace(module)

# High-level intermediate representation (IR) - Graph representation
print(symbolic_traced.graph)
"""
graph():
    %x : [num_users=1] = placeholder[target=x]
    %param : [num_users=1] = get_attr[target=param]
    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
    %linear : [num_users=1] = call_module[target=linear](args = (%add,), kwargs = {})
    %clamp : [num_users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
    return clamp
"""

# Code generation - valid Python code
print(symbolic_traced.code)
"""
def forward(self, x):
    param = self.param
    add = x + param;  x = param = None
    linear = self.linear(add);  add = None
    clamp = linear.clamp(min = 0.0, max = 1.0);  linear = None
    return clamp
"""

符号跟踪器 对 Python 代码执行“符号执行”。它将假值（称为代理 Proxies）传递给代码。对这些代理的操作会被记录下来。有关符号跟踪的更多信息可以在 symbolic_trace() 和 Tracer 文档中找到。

中间表示（Intermediate Representation，简称IR）是符号跟踪期间记录的操作容器。它由一系列节点组成，这些节点代表函数输入、调用站点（包括函数、方法或torch.nn.Module 实例），以及返回值。关于IR的更多信息可以在 Graph 文档中找到。IR是应用变换的基础格式。

Python代码生成 是使FX成为Python到Python（或模块到模块）转换工具包的关键。对于每个图IR，我们可以创建与其语义匹配的有效Python代码。此功能封装在GraphModule中，它是持有Graph以及从该图生成的forward方法的torch.nn.Module实例。

这些组件组成的管道（符号跟踪 -> 中间表示 -> 转换 -> Python 代码生成）共同构成了 FX 的 Python 到 Python 转换管道。此外，这些组件可以单独使用。例如，符号跟踪可以在隔离状态下用于捕获代码的一种形式以进行分析（而非转换）。代码生成可用于根据配置文件等程序化地生成模型。FX 具有许多用途！

几个示例转换可以在示例仓库中找到。

编写转换

什么是FX变换？本质上，它是一个这样的函数。

import torch
import torch.fx

def transform(m: nn.Module,
              tracer_class : type = torch.fx.Tracer) -> torch.nn.Module:
    # Step 1: Acquire a Graph representing the code in `m`

    # NOTE: torch.fx.symbolic_trace is a wrapper around a call to
    # fx.Tracer.trace and constructing a GraphModule. We'll
    # split that out in our transform to allow the caller to
    # customize tracing behavior.
    graph : torch.fx.Graph = tracer_class().trace(m)

    # Step 2: Modify this Graph or create a new one
    graph = ...

    # Step 3: Construct a Module to return
    return torch.fx.GraphModule(m, graph)

你的转换会接收一个torch.nn.Module，从中获取一个Graph，进行一些修改，并返回一个新的torch.nn.Module。你应该将你的FX转换返回的torch.nn.Module视为与常规的torch.nn.Module相同——你可以将其传递给另一个FX转换，可以传递给TorchScript，或者直接运行它。确保你的FX转换的输入和输出是一个torch.nn.Module将使其具有组合性。

import torch
import torch.fx

def transform(m : nn.Module) -> nn.Module:
    gm : torch.fx.GraphModule = torch.fx.symbolic_trace(m)

    # Modify gm.graph
    # <...>

    # Recompile the forward() method of `gm` from its Graph
    gm.recompile()

    return gm

鉴于你传入了一个已经被追踪为 torch.nn.Module 的 Graph，现在有两种主要的方法可以用来构建一个新的 Graph。

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)

gm.graph.print_tabular()

import torch
import torch.fx

# Sample module
class M(torch.nn.Module):
    def forward(self, x, y):
        return torch.add(x, y)

def transform(m: torch.nn.Module,
              tracer_class : type = fx.Tracer) -> torch.nn.Module:
    graph : fx.Graph = tracer_class().trace(m)
    # FX represents its Graph as an ordered list of
    # nodes, so we can iterate through them.
    for node in graph.nodes:
        # Checks if we're calling a function (i.e:
        # torch.add)
        if node.op == 'call_function':
            # The target attribute is the function
            # that call_function calls.
            if node.target == torch.add:
                node.target = torch.mul

    graph.lint() # Does some checks to make sure the
                 # Graph is well-formed.

    return fx.GraphModule(m, graph)

# Specifies the insertion point. Any nodes added to the
# Graph within this scope will be inserted after `node`
with traced.graph.inserting_after(node):
    # Insert a new `call_function` node calling `torch.relu`
    new_node = traced.graph.call_function(
        torch.relu, args=(node,))

    # We want all places that used the value of `node` to
    # now use that value after the `relu` call we've added.
    # We use the `replace_all_uses_with` API to do this.
    node.replace_all_uses_with(new_node)

# Note that this decomposition rule can be read as regular Python
def relu_decomposition(x):
    return (x > 0) * x

decomposition_rules = {}
decomposition_rules[F.relu] = relu_decomposition

def decompose(model: torch.nn.Module,
              tracer_class : type = fx.Tracer) -> torch.nn.Module:
"""
    Decompose `model` into smaller constituent operations.
    Currently,this only supports decomposing ReLU into its
    mathematical definition: (x > 0) * x
    """
    graph : fx.Graph = tracer_class().trace(model)
    new_graph = fx.Graph()
    env = {}
    tracer = torch.fx.proxy.GraphAppendingTracer(new_graph)
    for node in graph.nodes:
        if node.op == 'call_function' and node.target in decomposition_rules:
            # By wrapping the arguments with proxies,
            # we can dispatch to the appropriate
            # decomposition rule and implicitly add it
            # to the Graph by symbolically tracing it.
            proxy_args = [
                fx.Proxy(env[x.name], tracer) if isinstance(x, fx.Node) else x for x in node.args]
            output_proxy = decomposition_rules[node.target](*proxy_args)

            # Operations on `Proxy` always yield new `Proxy`s, and the
            # return value of our decomposition rule is no exception.
            # We need to extract the underlying `Node` from the `Proxy`
            # to use it in subsequent iterations of this transform.
            new_node = output_proxy.node
            env[node.name] = new_node
        else:
            # Default case: we don't have a decomposition rule for this
            # node, so just copy the node over into the new graph.
            new_node = new_graph.node_copy(node, lambda x: env[x.name])
            env[node.name] = new_node
    return fx.GraphModule(model, new_graph)

import torch
import torch.fx
from torch.fx.node import Node

from typing import Dict

class ShapeProp:
"""
    Shape propagation. This class takes a `GraphModule`.
    Then, its `propagate` method executes the `GraphModule`
    node-by-node with the given arguments. As each operation
    executes, the ShapeProp class stores away the shape and
    element type for the output values of each operation on
    the `shape` and `dtype` attributes of the operation's
    `Node`.
    """
    def __init__(self, mod):
        self.mod = mod
        self.graph = mod.graph
        self.modules = dict(self.mod.named_modules())

    def propagate(self, *args):
        args_iter = iter(args)
        env : Dict[str, Node] = {}

        def load_arg(a):
            return torch.fx.graph.map_arg(a, lambda n: env[n.name])

        def fetch_attr(target : str):
            target_atoms = target.split('.')
            attr_itr = self.mod
            for i, atom in enumerate(target_atoms):
                if not hasattr(attr_itr, atom):
                    raise RuntimeError(f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}")
                attr_itr = getattr(attr_itr, atom)
            return attr_itr

        for node in self.graph.nodes:
            if node.op == 'placeholder':
                result = next(args_iter)
            elif node.op == 'get_attr':
                result = fetch_attr(node.target)
            elif node.op == 'call_function':
                result = node.target(*load_arg(node.args), **load_arg(node.kwargs))
            elif node.op == 'call_method':
                self_obj, *args = load_arg(node.args)
                kwargs = load_arg(node.kwargs)
                result = getattr(self_obj, node.target)(*args, **kwargs)
            elif node.op == 'call_module':
                result = self.modules[node.target](*load_arg(node.args), **load_arg(node.kwargs))

            # This is the only code specific to shape propagation.
            # you can delete this `if` branch and this becomes
            # a generic GraphModule interpreter.
            if isinstance(result, torch.Tensor):
                node.shape = result.shape
                node.dtype = result.dtype

            env[node.name] = result

        return load_arg(self.graph.result)

调试

import torch
import torch.fx
import torchvision.models as models

def transform(m : torch.nn.Module) -> torch.nn.Module:
    gm = torch.fx.symbolic_trace(m)

    # Imagine we're doing some transforms here
    # <...>

    gm.recompile()

    return gm

resnet18 = models.resnet18()
transformed_resnet18 = transform(resnet18)

input_image = torch.randn(5, 3, 224, 224)

assert resnet18(input_image) == transformed_resnet18(input_image)
"""
RuntimeError: Boolean value of Tensor with more than one value is ambiguous
"""

assert torch.allclose(resnet18(input_image), transformed_resnet18(input_image))

import torch
import torch.fx
import torchvision.models as models

def my_pass(inp: torch.nn.Module, tracer_class : type = fx.Tracer) -> torch.nn.Module:
    graph = tracer_class().trace(inp)
    # Transformation logic here
    # <...>

    # Return new Module
    return fx.GraphModule(inp, graph)

my_module = models.resnet18()
my_module_transformed = my_pass(my_module)

input_value = torch.randn(5, 3, 224, 224)

# When this line is executed at runtime, we will be dropped into an
# interactive `pdb` prompt. We can use the `step` or `s` command to
# step into the execution of the next line
import pdb; pdb.set_trace()

my_module_transformed(input_value)

# Assume that `traced` is a GraphModule that has undergone some
# number of transforms

# Copy this code for later
print(traced)
# Print the code generated from symbolic tracing. This outputs:
"""
def forward(self, y):
    x = self.x
    add_1 = x + y;  x = y = None
    return add_1
"""

# Subclass the original Module
class SubclassM(M):
    def __init__(self):
        super().__init__()

    # Paste the generated `forward` function (the one we printed and
    # copied above) here
    def forward(self, y):
        x = self.x
        add_1 = x + y;  x = y = None
        return add_1

# Create an instance of the original, untraced Module. Then, create an
# instance of the Module with the copied `forward` function. We can
# now compare the output of both the original and the traced version.
pre_trace = M()
post_trace = SubclassM()

m = symbolic_trace(M())
m.to_folder("foo", "Bar")
from foo import Bar
y = Bar()

# Sample Module
class M(torch.nn.Module):
    def forward(self, x, y):
        return x + y

# Create an instance of `M`
m = M()

# Symbolically trace an instance of `M` (returns a GraphModule). In
# this example, we'll only be discussing how to inspect a
# GraphModule, so we aren't showing any sample transforms for the
# sake of brevity.
traced = symbolic_trace(m)

# Print the code produced by tracing the module.
print(traced)
# The generated `forward` function is:
"""
def forward(self, x, y):
    add = x + y;  x = y = None
    return add
"""

# Print the internal Graph.
print(traced.graph)
# This print-out returns:
"""
graph():
    %x : [num_users=1] = placeholder[target=x]
    %y : [num_users=1] = placeholder[target=y]
    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %y), kwargs = {})
    return add
"""

# Print a tabular representation of the internal Graph.
traced.graph.print_tabular()
# This gives us:
"""
opcode         name    target                   args    kwargs
-------------  ------  -----------------------  ------  --------
placeholder    x       x                        ()      {}
placeholder    y       y                        ()      {}
call_function  add     <built-in function add>  (x, y)  {}
output         output  output                   (add,)  {}
"""

# Sample user-defined function
def transform_graph(module: torch.nn.Module, tracer_class : type = fx.Tracer) -> torch.nn.Module:
    # Get the Graph from our traced Module
    g = tracer_class().trace(module)

"""
    Transformations on `g` go here
    """

    return fx.GraphModule(module, g)

# Transform the Graph
transformed = transform_graph(traced)

# Print the new code after our transforms. Check to see if it was
# what we expected
print(transformed)

符号跟踪的局限性

FX 使用一种称为 符号跟踪 （又称 符号执行）的系统，以可转换和分析的形式捕获程序语义。该系统通过执行程序（实际上是torch.nn.Module或函数）并记录操作来进行跟踪。在执行过程中，流经程序的数据不是实际数据，而是符号（在FX中称为 Proxy），因此它是符号化的。

虽然符号跟踪适用于大多数神经网络代码，但它也存在一些局限性。

def func_to_trace(x):
    if x.sum() > 0:
        return torch.relu(x)
    else:
        return torch.neg(x)

traced = torch.fx.symbolic_trace(func_to_trace)
"""
  <...>
  File "dyn.py", line 6, in func_to_trace
    if x.sum() > 0:
  File "pytorch/torch/fx/proxy.py", line 155, in __bool__
    return self.tracer.to_bool(self)
  File "pytorch/torch/fx/proxy.py", line 85, in to_bool
    raise TraceError('symbolically traced variables cannot be used as inputs to control flow')
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow
"""

import torch
import torch.fx

class MyModule(torch.nn.Module):
    def __init__(self, do_activation : bool = False):
        super().__init__()
        self.do_activation = do_activation
        self.linear = torch.nn.Linear(512, 512)

    def forward(self, x):
        x = self.linear(x)
        # This if-statement is so-called static control flow.
        # Its condition does not depend on any input values
        if self.do_activation:
            x = torch.relu(x)
        return x

without_activation = MyModule(do_activation=False)
with_activation = MyModule(do_activation=True)

traced_without_activation = torch.fx.symbolic_trace(without_activation)
print(traced_without_activation.code)
"""
def forward(self, x):
    linear_1 = self.linear(x);  x = None
    return linear_1
"""

traced_with_activation = torch.fx.symbolic_trace(with_activation)
print(traced_with_activation.code)
"""
import torch
def forward(self, x):
    linear_1 = self.linear(x);  x = None
    relu_1 = torch.relu(linear_1);  linear_1 = None
    return relu_1
"""

def f(x, flag):
    if flag: return x
    else: return x*2

fx.symbolic_trace(f) # Fails!

fx.symbolic_trace(f, concrete_args={'flag': True})

import torch
import torch.fx
from math import sqrt

def normalize(x):
"""
    Normalize `x` by the size of the batch dimension
    """
    return x / sqrt(len(x))

# It's valid Python code
normalize(torch.rand(3, 4))

traced = torch.fx.symbolic_trace(normalize)
"""
  <...>
  File "sqrt.py", line 9, in normalize
    return x / sqrt(len(x))
  File "pytorch/torch/fx/proxy.py", line 161, in __len__
    raise RuntimeError("'len' is not supported in symbolic tracing by default. If you want "
RuntimeError: 'len' is not supported in symbolic tracing by default. If you want this call to be recorded, please call torch.fx.wrap('len') at module scope
"""

torch.fx.wrap('len')
torch.fx.wrap('sqrt')

traced = torch.fx.symbolic_trace(normalize)

print(traced.code)
"""
import math
def forward(self, x):
    len_1 = len(x)
    sqrt_1 = math.sqrt(len_1);  len_1 = None
    truediv = x / sqrt_1;  x = sqrt_1 = None
    return truediv
"""

class MyCustomTracer(torch.fx.Tracer):
    # Inside here you can override various methods
    # to customize tracing. See the `Tracer` API
    # reference
    pass


# Let's use this custom tracer to trace through this module
class MyModule(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x) + torch.ones(3, 4)

mod = MyModule()

traced_graph = MyCustomTracer().trace(mod)
# trace() returns a Graph. Let's wrap it up in a
# GraphModule to make it runnable
traced = torch.fx.GraphModule(mod, traced_graph)

class MySpecialSubmodule(torch.nn.Module):
    def forward(self, x):
        return torch.neg(x)

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 4)
        self.submod = MySpecialSubmodule()

    def forward(self, x):
        return self.submod(self.linear(x))

traced = torch.fx.symbolic_trace(MyModule())
print(traced.code)
# `linear` is preserved as a call, yet `submod` is traced though.
# This is because the default set of "Leaf Modules" includes all
# standard `torch.nn` modules.
"""
import torch
def forward(self, x):
    linear_1 = self.linear(x);  x = None
    neg_1 = torch.neg(linear_1);  linear_1 = None
    return neg_1
"""

API参考

torch.fx.symbolic_trace(root, concrete_args=None)[源代码]

符号追踪API

给定一个 nn.Module 或函数实例 root，此函数将返回一个通过在跟踪 root 时记录的操作构建的 GraphModule。

concrete_args 允许你部分地专门化你的函数，以移除控制流或数据结构。

例如：

def f(a, b):
    if b == True:
        return a
    else:
        return a*2

由于存在控制流，FX 通常无法进行追踪。但是，我们可以使用 concrete_args 来专门针对 b 的值进行追踪。

f = fx.symbolic_trace(f, concrete_args={'b': False})
assert f(3, False)  == 6

请注意，尽管你可以传入不同的b值，但这些值将会被忽略。

我们还可以使用concrete_args来移除函数中的数据结构处理。这会利用pytrees展开输入参数。为了避免过度专门化，对于不需要专门化的值应传递fx.PH。例如：

def f(x):
    out = 0
    for v in x.values():
        out += v
    return out
f = fx.symbolic_trace(f, concrete_args={'x': {'a': fx.PH, 'b': fx.PH, 'c': fx.PH}})
assert f({'a': 1, 'b': 2, 'c': 4}) == 7

参数

• root (Union[torch.nn.Module, Callable]) – 需要被追踪并转换为图表示的模块或函数。

• concrete_args (Optional[Dict[str, any]]) – 需要部分特化的输入参数

返回值

由root记录的操作生成的模块。

返回类型

GraphModule

注意

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

torch.fx.wrap(fn_or_name)[源代码]

此函数可以在模块级别的作用域中调用，用于将 fn_or_name 注册为“叶子函数”。一个“叶子函数”在 FX 跟踪中会保留为 CallFunction 节点，而不会被进一步追踪。

# foo/bar/baz.py
def my_custom_function(x, y):
    return x * x + y * y

torch.fx.wrap('my_custom_function')

def fn_to_be_traced(x, y):
    # When symbolic tracing, the below call to my_custom_function will be inserted into
    # the graph rather than tracing it.
    return my_custom_function(x, y)

此函数还可以等效地用作装饰器：

# foo/bar/baz.py
@torch.fx.wrap
def my_custom_function(x, y):
    return x * x + y * y

一个被包装的函数可以视为一个“叶函数”，类似于“叶模块”的概念。也就是说，在FX跟踪中，这些函数会被保留在调用层面，而不进行内部追踪。

参数

fn_or_name (Union[str, Callable]) – 在调用时插入图中的函数或全局函数的名称

注意

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

类torch.fx.GraphModule(*args, **kwargs)[源代码]

GraphModule 是由 fx.Graph 生成的一个 nn.Module。它拥有一个 graph 属性，以及从该 graph 生成的 code 和 forward 属性。

警告

当graph被重新赋值时，code和forward将自动再生。然而，如果你在不重新赋值graph属性的情况下编辑了graph的内容，你需要调用recompile()来更新生成的代码。

注意

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

__init__(root, graph, class_name='GraphModule')[源代码]

创建一个GraphModule。

参数

• root (Union[torch.nn.Module, Dict[str, Any]) – root 可以是一个 nn.Module 实例，也可以是将字符串映射到任意属性类型的字典。如果 root 是一个 Module，则图中的节点的 target 字段中引用的基于 Module 的对象（通过全限定名）将会从 root 的 Module 层次结构复制到 GraphModule 的模块层次结构。如果 root 是一个字典，则节点中的 target 中的全限定名将直接在字典中查找，该对象将会被复制到 GraphModule 模块层次结构中的适当位置。

• graph (Graph) – graph 包含此 GraphModule 用于代码生成的节点

• class_name (str) – name 表示此 GraphModule 的调试名称。如果未设置，所有错误消息将默认显示为来自 GraphModule。将其设置为 root 的原始名称或在转换上下文中具有意义的名称可能会有所帮助。

注意

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

add_submodule(target, m)[源代码]

将给定的子模块添加到 self。

如果它们是 target 的子路径且尚不存在，这会安装空的模块。

参数

• target (str) – 新子模块的完整qualified字符串名称（参见nn.Module.get_submodule中的示例，了解如何指定完整qualified字符串。）

• m (Module) – 当前 Module 中要安装的子模块对象

返回值

子模块是否可以被插入。

为了使该方法返回 True，target 标注的链中每个对象必须要么尚未存在，要么引用一个 nn.Module（而不是参数或其它属性）。

返回类型

bool

注意

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

属性 代码: str

返回由该 GraphModule 底层的 Graph 生成的 Python 代码。

delete_all_unused_submodules()[源代码]

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

如果满足以下任一条件，则认为一个模块被“使用”了：1. 它有被使用的子模块 2. 其 forward 方法通过一个 call_module 节点直接调用 3. 它有一个非模块属性，该属性从 get_attr 节点中被使用

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

注意

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

delete_submodule(target)[源代码]

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

如果 target 不是有效的目标，模块将不会被删除。

参数

target (str) – 新子模块的完整qualified字符串名称（参见nn.Module.get_submodule中的示例，了解如何指定完整qualified字符串。）

返回值

目标字符串是否引用了某项内容

这个返回值表示我们要删除的是否是一个有效的子模块引用。如果返回值是False，则说明target不是一个有效的子模块引用。

返回类型

bool

注意

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

property graph: Graph

返回此 GraphModule 的底层 Graph

print_readable(print_output=True, include_stride=False, include_device=False, colored=False)[源代码]

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

警告

此 API 是实验性的，且不支持向下兼容。

recompile()[源代码]

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

注意

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

返回类型

Python代码

to_folder(folder, module_name='FxModule')[源代码]

将模块输出到指定的 folder 文件夹中，并使用 module_name 进行命名，以便于后续操作。

使用from <folder> import <module_name>进行导入

参数：

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

module_name (str): 指定的顶级名称，用于标识 Module

写代码

警告

此 API 是实验性的，且不支持向下兼容。

类torch.fx.Graph(owning_module=None, tracer_cls=None, tracer_extras=None)[源代码]

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

创建一个空的图。

注意

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

call_function(the_function, args=None, kwargs=None, type_expr=None)[源代码]

在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)[源代码]

在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 包含一个名为 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) – 该节点的指令码。可以是 ‘call_function’, ‘call_method’, ‘get_attr’, ‘call_module’, ‘placeholder’ 或 ‘output’ 中的一个。这些指令码的具体含义在 Graph 的文档字符串中进行了描述。

• args (Optional[Tuple[Argument, ...]]) – 此节点的参数以元组形式表示。

• kwargs (Optional[Dict[str, Argument]]) – 此节点的 kwargs 参数

• name (可选[str]) – 可选的字符串名称，用于标识Node。这将影响在生成的 Python 代码中该值的变量名。

• type_expr (Optional[Any]) – 表示该节点输出的 Python 类型的可选类型注解。

返回值

新创建并插入的节点。

返回类型

节点

注意

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

eliminate_dead_code(is_impure_node=None)[源代码]

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

参数

• is_impure_node (Optional[Callable[ [Node] , bool]] ) – 返回一个函数

• None (（节点是否为不纯。如果是）) –

• 到 （即默认行为）–

• Node.is_impure （使用）–

返回值

该pass是否改变了图。

返回类型

bool

示例：

在删除死代码之前，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 中删除的节点。

注意

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

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

支持快速查询节点

参数

• op (str) – 操作的名字

• target (Optional[Target]) - 节点的目标。在 call_function 中，目标是必需的；而对于其他操作，目标则是可选的。

• sort (bool) - 是否按照节点在图中出现的顺序来返回节点。

返回值

具有指定操作和目标的节点集合的可迭代对象。

警告

此 API 是实验性的，且不支持向下兼容。

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

在图中插入一个get_attr节点。这个get_attr节点表示从Module层次结构中获取属性。

参数

• qualified_name (str) – 要检索的属性的完全限定名称。例如，如果被跟踪模块包含一个名为foo的子模块，该子模块又包含一个名为bar的子模块，并且这个子模块有一个名为baz的属性，则应将完全限定名称foo.bar.baz作为参数传递。

• 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 节点，那么现在与 g 的输出值等价的 self 中的值。否则为 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)[源代码]

设置创建节点及其相关方法在图中的插入位置。

当在“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)[源代码]

设置创建节点及其相关方法在图中的插入位置。

当在“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()[源代码]

对此图进行各种检查，确保其结构正确。具体包括：
- 确认节点的拥有权是否正确（由该图所有）
- 验证节点是否按拓扑顺序排列
- 如果此图属于一个GraphModule，则检查目标是否存在于此GraphModule中

注意

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

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

将一个节点从一个图复制到另一个图中。arg_transform 需要将源节点图中的参数转换为目标图中的相应参数。示例：

# 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 转换为在 self 中等效的参数。 在最简单的情况下，这个函数会从映射原始图中的 Node 到 self 的表中检索值。

返回类型

节点

注意

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

属性节点: _node_list

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

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

返回值

一个由节点组成的双向链表。注意，可以通过调用reversed来切换迭代顺序。

on_generate_code(make_transformer)[源代码]

在生成 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 是实验性的，且不支持向下兼容。

output(result, type_expr=None)[源代码]

在Graph中插入一个output Node。一个output节点表示Python代码中的return语句。result是要返回的结果。

参数

• result (参数) - 需要返回的值。

• type_expr (Optional[Any]) – 表示该节点输出的 Python 类型的可选类型注解。

注意

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

注意

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

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

在图中插入一个占位符节点，它表示函数的输入。

参数

• name (str) – 输入值的名称。这对应于该 Graph 表示的函数的位置参数。

• type_expr (Optional[Any]) – 可选的类型注解，表示此节点输出的 Python 类型。在某些情况下（例如，在 TorchScript 编译过程中使用该函数时），这对于正确的代码生成是必需的。

• default_value (Any) – 该函数参数的默认值。注意：为了允许 None 作为默认值，应将 inspect.Signature.empty 传递给此参数以指定参数没有默认值。

返回类型

节点

注意

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

注意

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

print_tabular()[源代码]

以表格形式打印图的中间表示。请注意，使用此功能需要先安装 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)[源代码]

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

参数

root_module (str) – 根模块的名称，用于查找qualified name目标。这通常是‘self’。

返回值

src: 表示对象的 Python 源代码；globals: 包含 src 中全局名称及其引用对象的字典。

返回类型

一个PythonCode对象，包含两个字段。

注意

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

set_codegen(codegen)[源代码]

警告

此 API 是实验性的，且不支持向下兼容。

classtorch.fx.Node(graph, name, op, target, args, kwargs, return_type=None)[源代码]

Node 是表示 Graph 中各个操作的数据结构。通常情况下，Node 表示对各种实体（如运算符、方法和模块）的调用位置。一些例外情况包括指定函数输入和输出的节点。每个 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 的向后兼容性得到了保证。

属性all_input_nodes:List[Node]

返回所有输入到该节点的节点。这相当于遍历 args 和 kwargs，并只收集其中为节点的值。

返回值

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

append(x)[源代码]

在此节点的节点列表中插入 x。相当于执行 self.next.prepend(x)。

参数

x (Node) – 需要放置在此节点后的节点。此节点和目标节点必须属于同一张图。

注意

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

propertyargs: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)[源代码]

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

此方法可以在不带任何参数的情况下用作调试工具。

此函数还用于 Graph 类的内部 __str__ 方法中。在该图周围 GraphModule 中自动生成的 forward 函数签名由 placeholder_names 和 maybe_return_typename 中的字符串组成。placeholder_names 和 maybe_return_typename 不应以其他方式使用。

参数

• placeholder_names (Optional[List[str]]) – 一个存储生成的 forward 函数中占位符格式化字符串的列表。仅供内部使用。

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

返回值

如果 1) 我们使用 format_node 作为内部辅助函数

在Graph的__str__方法中，如果self是一个占位符节点，则返回None。否则，返回当前节点的描述性字符串表示。

返回类型

str

注意

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

insert_arg(idx, arg)[源代码]

在指定索引的位置向参数列表中插入一个参数。

参数

• idx (int) – 在 self.args 中要插入的元素之前的索引位置。

• arg (参数) – 需要插入到 args 中的新参数值

注意

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

is_impure()[源代码]

返回该操作是否为不纯操作，即该操作是一个占位符或输出，或者是不纯的 call_function 或 call_module。

返回值

如果是不纯的操作或不是纯函数。

返回类型

bool

警告

此 API 是实验性的，且不支持向下兼容。

propertykwargs: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的文档字符串。

允许对此属性进行赋值。在赋值时，所有使用情况和用户的相关记录会自动更新。

属性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 为 true 时按位置顺序返回仅 kwargs 参数。还会填充默认值。不支持仅限位置的参数和可变参数。

支持模块调用。

可能需要arg_types和kwarg_types来区分函数的重载。

参数

• root (torch.nn.Module) – 用于解析模块目标的基础模块。

• arg_types (Optional[Tuple[Any]]) – 参数类型的元组

• kwarg_types (Optional[Dict[str, Any]]) – kwargs 的类型参数字典

• normalize_to_only_use_kwargs (bool) – 是否将规范改为只使用关键字参数。

返回值

返回 NamedTuple ArgsKwargsPair，如果失败则返回 None。

返回类型

Optional[ArgsKwargsPair]

警告

此 API 是实验性的，且不支持向下兼容。

prepend(x)[源代码]

在此节点前面的图的节点列表中插入 x。 示例：

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

参数

x (Node) – 需要放置在当前节点之前的一个节点，且必须属于同一张图。

注意

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

属性prev: Node

返回链表中之前的 Node 节点。

返回值

前一个 Node，在节点的链表中。

replace_all_uses_with(replace_with, delete_user_cb=<function Node.<lambda>>, *, propagate_meta=False)[源代码]

将图中所有的 self 替换为节点 replace_with。

参数

• replace_with (Node) - 用于替换所有 self 使用的节点。

• delete_user_cb (Callable) – 一个回调函数，用于决定是否应该删除自节点中的特定用户。

• propagate_meta (bool) – 是否将原始节点的.meta字段上的所有属性复制到替换节点上。出于安全考虑，此操作仅在替换节点尚未具有现有的.meta字段时才有效。

返回值

在此更改中涉及的节点列表。

返回类型

Node

注意

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

replace_input_with(old_input, new_input)[源代码]

遍历 self 的输入节点，将其中所有的 old_input 替换为 new_input。

参数

• old_input (Node) – 需要被替换的旧输入节点。

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

注意

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

属性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)[源代码]

将现有位置参数的值更新为 arg。更新后，self.args[idx] == arg。

参数

• idx (int) – 在 self.args 中要更新的元素的索引

• arg (参数) – 新的参数值，将被写入 args

注意

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

update_kwarg(key, arg)[源代码]

将现有关键字参数的值更新为 arg。更新后，self.kwargs[key] == arg。

参数

• key (str) – 需要更新的元素在 self.kwargs 中的键

• arg (参数) – 需要写入 kwargs 的新的参数值

注意

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