diff --git a/ngraph/frontend/onnx_editor/src/detail/subgraph_extraction.cpp b/ngraph/frontend/onnx_editor/src/detail/subgraph_extraction.cpp index 079c566fd2e..9bd2692e082 100644 --- a/ngraph/frontend/onnx_editor/src/detail/subgraph_extraction.cpp +++ b/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}}; } } diff --git a/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_deeper_head_cut.prototxt b/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_deeper_head_cut.prototxt index 6c718ff0b39..bbd569818b2 100644 --- a/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_deeper_head_cut.prototxt +++ b/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 diff --git a/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_head_cut.prototxt b/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_head_cut.prototxt index 46c5056cb63..41984db10ea 100644 --- a/ngraph/test/models/onnx/model_editor/reference/subgraph__linear_model_head_cut.prototxt +++ b/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 diff --git a/ngraph/test/onnx/onnx_editor.cpp b/ngraph/test/onnx/onnx_editor.cpp index defcf6745ba..84aa9333b3c 100644 --- a/ngraph/test/onnx/onnx_editor.cpp +++ b/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{"conv1/7x7_s2_2:conv1/7x7_s2_1"})); + EXPECT_EQ(editor.model_inputs(), (std::vector{"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(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)); +}