[ONNX] Fix adding output nodes in NNCFGraph (#2538)

### Changes Do not add edges for the output node if it has no parents. E.g. output is put on the initializer. ### Reason for changes 125248 ### Related tickets 125248 ### Tests Checked on reproducer in the ticket
openvinotoolkit · Apr 18, 2024 · 96eef35 · 96eef35
1 parent 48c67bd
commit 96eef35
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 24 deletions.
diff --git a/nncf/onnx/graph/nncf_graph_builder.py b/nncf/onnx/graph/nncf_graph_builder.py
@@ -297,31 +297,33 @@ def _add_nncf_output_nodes(
         """
         for i, _output in enumerate(model.graph.output):
             output_name = _output.name
-            layer_attributes = ONNXLayerAttributes()
-            output_node = nncf_graph.add_nncf_node(
-                node_name=MODEL_OUTPUT_OP_NAME + "_" + str(i),
-                node_type=NNCFGraphNodeType.OUTPUT_NODE,
-                node_metatype=OutputNoopMetatype,
-                layer_attributes=layer_attributes,
-            )
-            from_node = parents_node_mapping[output_name]
+            if output_name in parents_node_mapping:
+                layer_attributes = ONNXLayerAttributes()
+                output_node = nncf_graph.add_nncf_node(
+                    node_name=MODEL_OUTPUT_OP_NAME + "_" + str(i),
+                    node_type=NNCFGraphNodeType.OUTPUT_NODE,
+                    node_metatype=OutputNoopMetatype,
+                    layer_attributes=layer_attributes,
+                )
 
-            output_node_node_id = output_node.node_id
-            edge = edge_info_mapping[output_name]
-            output_shape = get_edge_shape(edge)
-            onnx_dtype = get_edge_dtype(edge)
-            nncf_dtype = GraphConverter.convert_onnx_dtype_to_nncf_dtype(onnx_dtype)
-            input_port_id = 0
-            from_node_id = nncf_graph.get_node_by_name(from_node.name).node_id
-            output_port_id = get_output_port_id_for_node_before_output(output_name, from_node)
-            nncf_graph.add_edge_between_nncf_nodes(
-                from_node_id=from_node_id,
-                to_node_id=output_node_node_id,
-                tensor_shape=output_shape,
-                input_port_id=input_port_id,
-                output_port_id=output_port_id,
-                dtype=nncf_dtype,
-            )
+                from_node = parents_node_mapping[output_name]
+
+                output_node_node_id = output_node.node_id
+                edge = edge_info_mapping[output_name]
+                output_shape = get_edge_shape(edge)
+                onnx_dtype = get_edge_dtype(edge)
+                nncf_dtype = GraphConverter.convert_onnx_dtype_to_nncf_dtype(onnx_dtype)
+                input_port_id = 0
+                from_node_id = nncf_graph.get_node_by_name(from_node.name).node_id
+                output_port_id = get_output_port_id_for_node_before_output(output_name, from_node)
+                nncf_graph.add_edge_between_nncf_nodes(
+                    from_node_id=from_node_id,
+                    to_node_id=output_node_node_id,
+                    tensor_shape=output_shape,
+                    input_port_id=input_port_id,
+                    output_port_id=output_port_id,
+                    dtype=nncf_dtype,
+                )
 
     @staticmethod
     def convert_onnx_dtype_to_nncf_dtype(onnx_dtype: int) -> Dtype:

diff --git a/...onnx/data/reference_graphs/original_nncf_graph/synthetic/output_with_no_parents_model.dot b/...onnx/data/reference_graphs/original_nncf_graph/synthetic/output_with_no_parents_model.dot
@@ -0,0 +1,7 @@
+strict digraph  {
+"0 Conv1" [id=0, type=Conv];
+"1 nncf_model_input_0" [id=1, type=nncf_model_input];
+"2 nncf_model_output_0" [id=2, type=nncf_model_output];
+"0 Conv1" -> "2 nncf_model_output_0"  [label="[]", style=solid];
+"1 nncf_model_input_0" -> "0 Conv1"  [label="[1, 3, 10, 10]", style=solid];
+}
diff --git a/tests/onnx/test_nncf_graph_builder.py b/tests/onnx/test_nncf_graph_builder.py
@@ -16,9 +16,11 @@
 import torch
 from torchvision import models
 
+from nncf.onnx.graph.model_transformer import ONNXModelTransformer
 from nncf.onnx.graph.nncf_graph_builder import GraphConverter
 from tests.onnx.conftest import ONNX_TEST_ROOT
 from tests.onnx.models import ALL_SYNTHETIC_MODELS
+from tests.onnx.models import OneConvolutionalModel
 from tests.onnx.opset_converter import convert_opset_version
 from tests.onnx.quantization.common import ModelToTest
 from tests.onnx.weightless_model import load_model_topology_with_zeros_weights
@@ -99,3 +101,13 @@ def test_compare_nncf_graph_detection_real_models(tmp_path, model_to_test):
     nx_graph = nncf_graph.get_graph_for_structure_analysis(extended=True)
 
     compare_nx_graph_with_reference(nx_graph, path_to_dot, check_edge_attrs=True)
+
+
+def test_add_output_nodes_with_no_parents_node():
+    model_to_test = OneConvolutionalModel().onnx_model
+    model_outputs = (value_info.name for value_info in model_to_test.graph.output)
+    model_with_output = ONNXModelTransformer._insert_outputs(model_to_test, (*model_outputs, "Conv1_W"))
+    nncf_graph = GraphConverter.create_nncf_graph(model_with_output)
+    nx_graph = nncf_graph.get_graph_for_structure_analysis(extended=True)
+    path_to_dot = REFERENCE_GRAPHS_DIR / "synthetic" / "output_with_no_parents_model.dot"
+    compare_nx_graph_with_reference(nx_graph, path_to_dot, check_edge_attrs=True)