torch.fx.subgraph_rewriter.replace_pattern

torch.fx.subgraph_rewriter.replace_pattern(gm, pattern, replacement) [источник]

Находит все возможные непересекающиеся наборы операторов и их зависимостей по данным (pattern) в графе GraphModule (gm), а затем заменяет каждый найденный подграф другим подграфом (replacement).

Параметры:

• gm (GraphModule) – GraphModule, содержащий граф, над которым выполняется операция
• pattern (Callable[[...], Any] | GraphModule) – Подграф, который нужно найти в gm для замены
• replacement (Callable[[...], Any] | GraphModule) – Подграф, которым нужно заменить 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]

Тип возвращаемого значения:

List[Match]

Примеры:

import torch
from torch.fx import symbolic_trace, subgraph_rewriter


class M(torch.nn.Module):
    def __init__(self) -> None:
        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])


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


traced_module = symbolic_trace(M())

subgraph_rewriter.replace_pattern(traced_module, pattern, replacement)

Приведённый выше код сначала найдёт pattern в методе forward объекта traced_module. Сопоставление с образцом выполняется на основе отношений использования и определения, а не имён узлов. Например, если бы в pattern у вас был p = torch.cat([a, b]), вы могли бы найти m = torch.cat([a, b]) в исходной функции forward, несмотря на различие имён переменных (p и m).

Оператор return в pattern сопоставляется только по значению; он может совпасть или не совпасть с оператором return в большем графе. Иными словами, образец не обязан доходить до конца большего графа.

После сопоставления образец удаляется из большей функции и заменяется на replacement. Если в большей функции есть несколько совпадений для pattern, будут заменены все непересекающиеся совпадения. Если совпадения пересекаются, будет заменено первое найденное совпадение из набора пересекающихся совпадений. («Первое» здесь определяется как первое в топологическом порядке отношений использования и определения узлов. В большинстве случаев первым узлом является параметр, непосредственно следующий за self, а последним — узел, возвращаемый функцией.)

Важно отметить, что параметры Callable pattern должны использоваться в самом Callable, а параметры Callable replacement должны соответствовать образцу. Первое правило объясняет, почему в приведённом выше блоке кода функция 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), даже если параметр y не используется в replacement.

После вызова 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 гарантируется обратная совместимость.