From 192039841cdd257dff56fdc909ad480d6173005a Mon Sep 17 00:00:00 2001 From: Daniil Lyakhov Date: Thu, 28 Sep 2023 12:16:25 +0200 Subject: [PATCH] [NNCFGraph] Fix `get_input_edges` for parallel edges (#2107) ### Changes * `parallel_input_port_ids` parameter plugged in to `OVNNCFGraph` * `NNCFGraph.get_input_edges` unwraps `parallel_input_port_ids` and returns separate entities of `NNCFGraphEdge` for each parallel edge ### Reason for changes * To make possible to insert statistics collectors on parallel edges in OV backend ### Related tickets 119666 ### Tests * tests/common/graph/test_nncf_graph.py is updated * tests/openvino/native/test_nncf_graph_builder.py is updated --- nncf/common/graph/graph.py | 48 ++++++++++++--- .../common/tensor_statistics/collectors.py | 2 +- nncf/openvino/graph/nncf_graph_builder.py | 14 ++++- tests/common/graph/test_nncf_graph.py | 59 +++++++++++++++++++ tests/openvino/native/models.py | 13 ++++ .../native/test_nncf_graph_builder.py | 48 +++++++++++++++ tests/torch/test_graph_analysis.py | 6 +- tests/torch/test_graph_building.py | 55 +++++++++++++++-- 8 files changed, 226 insertions(+), 19 deletions(-) diff --git a/nncf/common/graph/graph.py b/nncf/common/graph/graph.py index 2d696c9ee2f..0b5b3cf3db4 100644 --- a/nncf/common/graph/graph.py +++ b/nncf/common/graph/graph.py @@ -148,17 +148,23 @@ def __init__( self.parallel_input_port_ids = parallel_input_port_ids def __str__(self): - return str(self.from_node) + " -> " + str(self.tensor_shape) + " -> " + str(self.to_node) + return f"{self.from_node}:{self.output_port_id} -> {self.tensor_shape} -> {self.to_node}:{self.input_port_id}" def __hash__(self): - return hash(str(self)) + return hash( + ( + self.from_node, + self.to_node, + self.input_port_id, + self.output_port_id, + tuple(self.tensor_shape), + self.dtype, + tuple(self.parallel_input_port_ids), + ) + ) def __eq__(self, other): - return ( - self.from_node == other.from_node - and self.to_node == other.to_node - and self.tensor_shape == other.tensor_shape - ) + return isinstance(other, NNCFGraphEdge) and self.__dict__ == other.__dict__ class NNCFGraphPatternIO: @@ -331,7 +337,9 @@ def get_input_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]: :return: List of input edges for the node sorted by input port ID. """ input_nodes = self.get_previous_nodes(node) - edges = [self.get_edge(from_node, node) for from_node in input_nodes] + edges = [] + for from_node in input_nodes: + edges.extend(self._get_edges(from_node, node)) return sorted(edges, key=lambda x: x.input_port_id) def get_output_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]: @@ -343,9 +351,31 @@ def get_output_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]: """ output_nodes = self.get_next_nodes(node) - edges = [self.get_edge(node, to_node) for to_node in output_nodes] + edges = [] + for to_node in output_nodes: + edges.extend(self._get_edges(node, to_node)) return sorted(edges, key=lambda x: x.output_port_id) + def _get_edges(self, from_node: NNCFNode, to_node: NNCFNode) -> List[NNCFGraphEdge]: + edges = [] + edge = self.get_edge(from_node, to_node) + parallel_input_port_ids = edge.parallel_input_port_ids + edge.parallel_input_port_ids = [] + edges.append(edge) + for input_port_id in parallel_input_port_ids: + edges.append( + NNCFGraphEdge( + from_node=edge.from_node, + to_node=edge.to_node, + input_port_id=input_port_id, + output_port_id=edge.output_port_id, + tensor_shape=edge.tensor_shape, + dtype=edge.dtype, + parallel_input_port_ids=[], + ) + ) + return edges + def traverse_graph( self, curr_node: NNCFNode, diff --git a/nncf/experimental/common/tensor_statistics/collectors.py b/nncf/experimental/common/tensor_statistics/collectors.py index 3655fffe5d6..08bf166f32d 100644 --- a/nncf/experimental/common/tensor_statistics/collectors.py +++ b/nncf/experimental/common/tensor_statistics/collectors.py @@ -366,7 +366,7 @@ def get_tensor_collector_inputs( :param outputs: Target model outputs. :param output_info: Output info collected by a `TensorCollector.get_output_info` method. - :returns: Model outputs in a format required by `TensorCollector.register_input` method. + :returns: Model outputs in a format required by `TensorCollector.register_inputs` method. """ target_inputs = {} for reducer, names in output_info: diff --git a/nncf/openvino/graph/nncf_graph_builder.py b/nncf/openvino/graph/nncf_graph_builder.py index b1101224127..1a5374c65e7 100644 --- a/nncf/openvino/graph/nncf_graph_builder.py +++ b/nncf/openvino/graph/nncf_graph_builder.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections import defaultdict from typing import List, Type import openvino.runtime as ov @@ -88,19 +89,28 @@ def _add_edges_to_nncf_graph(model: ov.Model, graph: NNCFGraph) -> None: for op in model.get_ops(): in_node_id = graph.get_node_by_name(op.get_friendly_name()).node_id for output_port_id, out in enumerate(op.outputs()): + node_vs_target_inputs = defaultdict(list) for inp in out.get_target_inputs(): - out_node = inp.get_node() + node_vs_target_inputs[inp.get_node()].append(inp) + + for out_node, inputs in node_vs_target_inputs.items(): tensor_shape = list(out.partial_shape.get_max_shape()) output_node_id = graph.get_node_by_name(out_node.get_friendly_name()).node_id ov_dtype = out.get_element_type().get_type_name() nncf_dtype = GraphConverter.convert_to_nncf_dtype(ov_dtype) + + parallel_inputs = None + if len(inputs) > 1: + parallel_inputs = [inp.get_index() for inp in inputs[1:]] + graph.add_edge_between_nncf_nodes( from_node_id=in_node_id, to_node_id=output_node_id, tensor_shape=tensor_shape, - input_port_id=inp.get_index(), + input_port_id=inputs[0].get_index(), output_port_id=output_port_id, dtype=Dtype(nncf_dtype), + parallel_input_port_ids=parallel_inputs, ) @staticmethod diff --git a/tests/common/graph/test_nncf_graph.py b/tests/common/graph/test_nncf_graph.py index ad166e1d2fe..a9a6d04b2de 100644 --- a/tests/common/graph/test_nncf_graph.py +++ b/tests/common/graph/test_nncf_graph.py @@ -10,6 +10,7 @@ # limitations under the License. from nncf.common.graph.graph import NNCFGraph +from nncf.common.graph.graph import NNCFGraphEdge from nncf.common.graph.layer_attributes import Dtype from nncf.common.graph.patterns import GraphPattern @@ -48,3 +49,61 @@ def test_find_matching_subgraphs(): continue assert len(match) == 2 assert match == nodes[:2] + + +def test_parallel_edges(): + def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id): + return NNCFGraphEdge( + from_node, + to_node, + input_port_id=input_port_id, + output_port_id=output_port_id, + parallel_input_port_ids=[], + tensor_shape=(1, 2, 3), + dtype="dummy", + ) + + nncf_graph = NNCFGraph() + nodes = [] + for node in "abc": + nodes.append(nncf_graph.add_nncf_node(node, f"type_{node}", f"metatype_{node}")) + + nncf_graph.add_edge_between_nncf_nodes( + nodes[0].node_id, + nodes[1].node_id, + input_port_id=0, + output_port_id=0, + parallel_input_port_ids=list(range(1, 5)), + tensor_shape=(1, 2, 3), + dtype="dummy", + ) + nncf_graph.add_edge_between_nncf_nodes( + nodes[0].node_id, + nodes[2].node_id, + input_port_id=10, + output_port_id=15, + parallel_input_port_ids=[], + tensor_shape=(1, 2, 3), + dtype="dummy", + ) + output_edges = nncf_graph.get_output_edges(nodes[0]) + input_edges = nncf_graph.get_input_edges(nodes[1]) + assert len(input_edges) == 5 + assert len(output_edges) == 6 + assert input_edges == output_edges[:-1] + for input_port_id, edge in enumerate(input_edges): + ref_edge = _get_default_nncf_graph_edge( + nodes[0], + nodes[1], + input_port_id=input_port_id, + output_port_id=0, + ) + assert ref_edge == edge + + ordinary_edge = _get_default_nncf_graph_edge( + nodes[0], + nodes[2], + input_port_id=10, + output_port_id=15, + ) + assert ordinary_edge == output_edges[-1] diff --git a/tests/openvino/native/models.py b/tests/openvino/native/models.py index 6440636ea50..6414fc9e7a4 100644 --- a/tests/openvino/native/models.py +++ b/tests/openvino/native/models.py @@ -650,6 +650,19 @@ def _create_ov_model(self): return model +class ParallelEdgesModel(OVReferenceModel): + def _create_ov_model(self) -> ov.Model: + input_shape = [1, 3, 3] + + input_1 = opset.parameter(input_shape, name="Input") + mm = opset.matmul(input_1, input_1, False, False, name="Mm") + add = opset.add(input_1, np.array(1.0, dtype=np.float32), name="Add") + result_0 = opset.result(mm, name="Result_mm") + result_1 = opset.result(add, name="Result_add") + model = ov.Model([result_0, result_1], [input_1]) + return model + + @SYNTHETIC_MODELS.register() class UnifiedEmbeddingModel(OVReferenceModel): def _create_ov_model(self): diff --git a/tests/openvino/native/test_nncf_graph_builder.py b/tests/openvino/native/test_nncf_graph_builder.py index da78b0d2dc1..03ea07e9cc6 100644 --- a/tests/openvino/native/test_nncf_graph_builder.py +++ b/tests/openvino/native/test_nncf_graph_builder.py @@ -12,9 +12,13 @@ import openvino.runtime as ov import pytest +from nncf.common.graph.graph import NNCFGraphEdge +from nncf.common.graph.layer_attributes import Dtype +from nncf.openvino.graph.nncf_graph_builder import GraphConverter from tests.openvino.conftest import OPENVINO_NATIVE_TEST_ROOT from tests.openvino.native.common import compare_nncf_graphs from tests.openvino.native.models import SYNTHETIC_MODELS +from tests.openvino.native.models import ParallelEdgesModel from tests.openvino.omz_helpers import convert_model from tests.openvino.omz_helpers import download_model @@ -47,3 +51,47 @@ def test_compare_nncf_graph_omz_models(tmp_path, model_name): path_to_dot = REFERENCE_GRAPHS_DIR / f"{model_name}.dot" compare_nncf_graphs(model, path_to_dot) + + +def test_parallel_edges(): + def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id): + return NNCFGraphEdge( + from_node, + to_node, + input_port_id=input_port_id, + output_port_id=output_port_id, + tensor_shape=[1, 3, 3], + dtype=Dtype.FLOAT, + parallel_input_port_ids=[], + ) + + model = ParallelEdgesModel().ov_model + nncf_graph = GraphConverter.create_nncf_graph(model) + input_node = nncf_graph.get_node_by_name("Input") + mm_node = nncf_graph.get_node_by_name("Mm") + add_node = nncf_graph.get_node_by_name("Add") + ref_input_edges = { + _get_default_nncf_graph_edge( + input_node, + mm_node, + input_port_id=0, + output_port_id=0, + ), + _get_default_nncf_graph_edge( + input_node, + mm_node, + input_port_id=1, + output_port_id=0, + ), + } + ref_output_edges = ref_input_edges.copy() + ref_output_edges.add( + _get_default_nncf_graph_edge( + input_node, + add_node, + input_port_id=0, + output_port_id=0, + ) + ) + assert set(nncf_graph.get_input_edges(mm_node)) == ref_input_edges + assert set(nncf_graph.get_output_edges(input_node)) == ref_output_edges diff --git a/tests/torch/test_graph_analysis.py b/tests/torch/test_graph_analysis.py index 6a3e1615b70..ddc322192c3 100644 --- a/tests/torch/test_graph_analysis.py +++ b/tests/torch/test_graph_analysis.py @@ -78,7 +78,7 @@ def get_node(name: NNCFNodeName): ["/C_0"], NNCFGraphPatternIO( input_edges=[make_mock_edge("/B_0", "/C_0", input_port_id=0, output_port_id=0)], - output_edges=[make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0)], + output_edges=[make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0)], ), ), ( @@ -86,7 +86,7 @@ def get_node(name: NNCFNodeName): NNCFGraphPatternIO( input_edges=[], output_edges=[ - make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0), + make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0), make_mock_edge("/A_0", "/D_0", input_port_id=0, output_port_id=1), ], ), @@ -95,7 +95,7 @@ def get_node(name: NNCFNodeName): ["/D_0"], NNCFGraphPatternIO( input_edges=[ - make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0), + make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0), make_mock_edge("/A_0", "/D_0", input_port_id=0, output_port_id=1), ], output_edges=[ diff --git a/tests/torch/test_graph_building.py b/tests/torch/test_graph_building.py index 11fa4a213f3..52152772a7b 100644 --- a/tests/torch/test_graph_building.py +++ b/tests/torch/test_graph_building.py @@ -21,6 +21,7 @@ from nncf.common.graph.definitions import MODEL_INPUT_OP_NAME from nncf.common.graph.definitions import MODEL_OUTPUT_OP_NAME from nncf.common.graph.definitions import NNCFGraphNodeType +from nncf.common.graph.graph import NNCFGraphEdge from nncf.common.graph.layer_attributes import GetItemLayerAttributes from nncf.common.graph.layer_attributes import MultipleInputLayerAttributes from nncf.common.graph.layer_attributes import MultipleOutputLayerAttributes @@ -407,10 +408,7 @@ def test_split_attributes(input_shape): graph = graph_builder.build_graph(model) chunk_nodes_with_attributes = { "ModelForTestWithSplit/chunk_0": {"chunks": 2, "axis": 1}, - "ModelForTestWithSplit/unbind_0": {"chunks": 2, "axis": 1} - # TODO: fix NNCFGraph tracing so valid reference below - # will be generated by NNCF - #'ModelForTestWithSplit/unbind_0': {'chunks': 20, 'axis': 1} + "ModelForTestWithSplit/unbind_0": {"chunks": 20, "axis": 1}, } for node in graph.get_all_nodes(): @@ -460,6 +458,55 @@ def test_getitem_attributes(input_shape): assert getitem_nodes_with_attributes[node.node_name] is None +class ParallelEdgesModel(nn.Module): + def forward(self, x): + mm_res = torch.mm(x, x) + return mm_res, x + mm_res + + +def test_parallel_edges_in_nncf_graph(): + def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id): + return NNCFGraphEdge( + from_node, + to_node, + input_port_id=input_port_id, + output_port_id=output_port_id, + tensor_shape=(3, 3), + dtype=Dtype.FLOAT, + parallel_input_port_ids=[], + ) + + input_shape = (3, 3) + model = ParallelEdgesModel() + input_info = ModelInputInfo(input_shape) + graph_builder = GraphBuilder( + create_dummy_forward_fn( + [ + input_info, + ], + with_input_tracing=True, + with_output_tracing=True, + ) + ) + + nncf_graph = graph_builder.build_graph(model) + + input_node = nncf_graph.get_node_by_name("/nncf_model_input_0") + mm_node = nncf_graph.get_node_by_name("ParallelEdgesModel/mm_0") + ref_input_edges = { + _get_default_nncf_graph_edge(input_node, mm_node, input_port_id=0, output_port_id=0), + _get_default_nncf_graph_edge(input_node, mm_node, input_port_id=1, output_port_id=0), + } + mm_node_input_edges = nncf_graph.get_input_edges(mm_node) + assert set(mm_node_input_edges) == ref_input_edges + ref_output_edges = ref_input_edges.copy() + + add_node = nncf_graph.get_node_by_name("ParallelEdgesModel/__add___0") + ref_output_edges.add(_get_default_nncf_graph_edge(input_node, add_node, input_port_id=0, output_port_id=0)) + input_node_output_edges = nncf_graph.get_output_edges(input_node) + assert set(input_node_output_edges) == ref_output_edges + + TEST_KEYWORD_1 = "keyword1" TEST_KEYWORD_2 = "keyword2" INPUT_INFO_CONFIG_VS_FORWARD_ARGS = [