SubgraphMatcherWithNameNodeMap

class torch.fx.passes.utils.matcher_with_name_node_map_utils.SubgraphMatcherWithNameNodeMap(pattern_gm, match_output=False, match_placeholder=False, remove_overlapping_matches=True, ignore_literals=False) [исходный код]

Расширяет SubgraphMatcher, позволяя получать узлы сопоставленного подграфа по имени узла,

для этого шаблон должен иметь определённый формат (возвращать дополнительный словарь в качестве выходных данных, где ключом является имя узла, а значением — узел графа шаблона; подробнее см. пример)

Отличие от SubgraphMatcher заключается в том, что при инициализации он принимает pattern_gm GraphModule в качестве входных данных, поскольку нам необходимо изменять граф (для чего требуется recompile GraphModule)

Пример:

def pattern(x, weight):
    conv = F.conv2d(x, weight)
    relu = F.relu(conv)
    return relu, {"conv": conv, "relu": relu}


def target_graph(x, weight):
    conv = F.conv2d(x, weight)
    relu = F.relu(conv)
    relu *= 2
    return relu


pattern_gm = export(pattern, example_inputs).module()
target_gm = export(target_graph, example_inputs).module()
matcher = SubgraphMatcherWithNameNodeMap(pattern_gm)
matches = matcher.match(target_gm)
for match in matches:
    match.name_node_map["conv"].meta["annotation"] = ...

Предупреждение

Этот API является экспериментальным и НЕ имеет обратной совместимости.

match(graph, node_name_match='') [исходный код]

В возвращённом InternalMatch будет заполнено поле name_node_map — отображение имени узла (str) на целевой узел, например {"conv": target_conv_node, "relu": target_relu_node}

Для этого граф шаблона должен возвращать дополнительные выходные данные, представляющие соответствие имени узла узлу. Например, вместо:

def pattern(...):
    ...
    return relu

следует сделать так:

def pattern(...):
    ...
    return relu, {"conv": conv, "relu": relu}

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

list[InternalMatch]