Preserve the input names when extracting subgraphs (#5207)

ngraph/frontend/onnx_editor/src/detail/subgraph_extraction.cpp

@@ -159,8 +159,6 @@ namespace
                      edge.m_node_idx,
                      ". Cannot append a new graph input to this node.");
 
-        const std::string new_input_name = target_node.output(0) + ":" + edge.m_tensor_name;
-
         // if an edge is connected to an initializer, the initializer is removed and substituted
         // with an input
         if (is_graph_initializer(graph, edge.m_tensor_name))
@@ -173,10 +171,10 @@ namespace
             auto& new_input = *(graph.add_input());
             // copy the intermediate tensor properties to the newly created input
             new_input.MergeFrom(find_tensor_descriptor(graph, edge.m_tensor_name));
-            *(new_input.mutable_name()) = new_input_name;
+            *(new_input.mutable_name()) = edge.m_tensor_name;
             // attach the new graph input to the target node's input
-            *target_input = new_input_name;
-            return {true, InputEdge{edge.m_node_idx, new_input_name}};
+            *target_input = edge.m_tensor_name;
+            return {true, InputEdge{edge.m_node_idx, edge.m_tensor_name}};
         }
     }
 

ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_deeper_head_cut.prototxt

@@ -4,7 +4,7 @@ doc_string: "This model contains the first few nodes of the ONNX Inception V1 mo
 graph {
   name: "Inception V1 fragment"
   node {
-    input: "pool1/3x3_s2_1:conv1/7x7_s2_2"
+    input: "conv1/7x7_s2_2"
     output: "pool1/3x3_s2_1"
     name: ""
     op_type: "MaxPool"
@@ -31,7 +31,7 @@ graph {
   }
 
   input {
-    name: "pool1/3x3_s2_1:conv1/7x7_s2_2"
+    name: "conv1/7x7_s2_2"
     type {
       tensor_type {
         elem_type: 1

ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_head_cut.prototxt

@@ -4,7 +4,7 @@ doc_string: "This model contains the first few nodes of the ONNX Inception V1 mo
 graph {
   name: "Inception V1 fragment"
   node {
-    input: "conv1/7x7_s2_2:conv1/7x7_s2_1"
+    input: "conv1/7x7_s2_1"
     output: "conv1/7x7_s2_2"
     name: ""
     op_type: "Relu"
@@ -37,7 +37,7 @@ graph {
   }
 
   input {
-    name: "conv1/7x7_s2_2:conv1/7x7_s2_1"
+    name: "conv1/7x7_s2_1"
     type {
       tensor_type {
         elem_type: 1

ngraph/test/onnx/onnx_editor.cpp

@@ -655,7 +655,7 @@ NGRAPH_TEST(onnx_editor, subgraph__inputs_getter)
 
     editor.cut_graph_fragment({{InputEdge(1, "conv1/7x7_s2_1")}}, {});
 
-    EXPECT_EQ(editor.model_inputs(), (std::vector<std::string>{"conv1/7x7_s2_2:conv1/7x7_s2_1"}));
+    EXPECT_EQ(editor.model_inputs(), (std::vector<std::string>{"conv1/7x7_s2_1"}));
 }
 
 using TestEngine = test::INTERPRETER_Engine;
@@ -771,3 +771,24 @@ NGRAPH_TEST(onnx_editor, values__append_two_initializers_mixed_types)
     test_case.add_expected_output<int16_t>(Shape{2, 2, 1}, {1, 4, 5, 8});
     test_case.run();
 }
+
+NGRAPH_TEST(onnx_editor, combined__cut_and_replace_shape)
+{
+    ONNXModelEditor editor{file_util::path_join(
+        SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.prototxt")};
+
+    const auto new_shape = PartialShape({1, 64, 112, 112});
+    editor.cut_graph_fragment({{InputEdge(1, "conv1/7x7_s2_1")}}, {});
+    editor.set_input_shapes({{"conv1/7x7_s2_1", new_shape}});
+
+    const auto ref_model = file_util::path_join(
+        SERIALIZED_ZOO, "onnx/model_editor/reference/subgraph__linear_model_head_cut.prototxt");
+
+    const auto result = compare_onnx_models(editor.model_string(), ref_model);
+
+    EXPECT_TRUE(result.is_ok) << result.error_message;
+
+    const auto graph_inputs = editor.get_function()->get_parameters();
+    EXPECT_TRUE(
+        find_input(graph_inputs, "conv1/7x7_s2_1")->get_partial_shape().same_scheme(new_shape));
+}