Python PyTorch replace_pattern用法及代码示例

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

用法:

torch.fx.replace_pattern(gm, pattern, replacement)

参数：

• gm-包装要操作的 Graph 的 GraphModule

• pattern-gm 中匹配的子图进行替换

• replacement-替换pattern 的子图

返回：

Match 对象列表，表示原始图中与 pattern 匹配的位置。如果没有匹配项，则列表为空。 Match 定义为：

class Match(NamedTuple):
    # Node from which the match was found
    anchor: Node
    # Maps nodes in the pattern subgraph to nodes in the larger graph
    nodes_map: Dict[Node, Node]

返回类型：

列表[匹配]

匹配 GraphModule ( gm ) 图中所有可能的非重叠运算符集及其数据依赖性 ( pattern )，然后将每个匹配的子图替换为另一个子图 ( replacement )。

例子：

import torch
from torch.fx import symbolic_trace, subgraph_rewriter

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, w1, w2):
        m1 = torch.cat([w1, w2]).sum()
        m2 = torch.cat([w1, w2]).sum()
        return x + torch.max(m1) + torch.max(m2)

def pattern(w1, w2):
    return torch.cat([w1, w2]).sum()

def replacement(w1, w2):
    return torch.stack([w1, w2])

traced_module = symbolic_trace(M())

subgraph_rewriter.replace_pattern(traced_module, pattern, replacement)

上面的代码将首先匹配 traced_module 的 forward 方法中的 pattern 。 Pattern-matching 是基于 use-def 关系完成的，而不是节点名称。例如，如果您在 pattern 中有 p = torch.cat([a, b]) ，则可以匹配原始 forward 函数中的 m = torch.cat([a, b])，尽管变量名称不同( p 与 m )。

pattern 中的return 语句仅根据其值进行匹配；它可能与较大图中的 return 语句匹配，也可能不匹配。换句话说，模式不必延伸到更大图的末尾。

当模式匹配时，它将从较大的函数中删除并替换为 replacement 。如果较大的函数中有多个pattern 匹配，则每个不重叠的匹配都将被替换。在匹配重叠的情况下，重叠匹配集中第一个找到的匹配将被替换。 (“First” 在这里被定义为节点的 use-def 关系的拓扑排序中的第一个。在大多数情况下，第一个 Node 是直接出现在 self 之后的参数，而最后一个 Node 是函数返回的任何内容.)

需要注意的重要一点是pattern Callable 的参数必须在Callable 本身中使用，并且replacement Callable 的参数必须与模式匹配。第一条规则是为什么在上面的代码块中，forward 函数有参数 x, w1, w2 ，但 pattern 函数只有参数 w1, w2 。 pattern 不使用 x ，因此不应将 x 指定为参数。作为第二条规则的示例，请考虑替换

def pattern(x, y):
    return torch.neg(x) + torch.relu(y)

和

def replacement(x, y):
    return torch.relu(x)

在这种情况下，replacement 需要与 pattern 相同数量的参数(x 和 y )，即使在 replacement 中没有使用参数 y。

调用 subgraph_rewriter.replace_pattern 后，生成的 Python 代码如下所示：

def forward(self, x, w1, w2):
    stack_1 = torch.stack([w1, w2])
    sum_1 = stack_1.sum()
    stack_2 = torch.stack([w1, w2])
    sum_2 = stack_2.sum()
    max_1 = torch.max(sum_1)
    add_1 = x + max_1
    max_2 = torch.max(sum_2)
    add_2 = add_1 + max_2
    return add_2

注意

保证此 API 的向后兼容性。