Handling of cut + reshape in MO using ONNX FE (#9102)

src/bindings/python/tests/test_frontend/test_frontend_onnx_editor.py

@@ -51,7 +51,7 @@ from openvino.frontend import FrontEndManager
 def create_test_onnx_models():
     models = {}
     # Input model 1
-    add = onnx.helper.make_node("Add", inputs=["in1", "in2"], outputs=["add_out"])
+    add = onnx.helper.make_node("Add", inputs=["in1", "in2"], outputs=["add_out"], name="onnx_add_op")
     split = onnx.helper.make_node("Split", inputs=["add_out"],
                                   outputs=["out1", "out2"], name="split1", axis=0)
     relu = onnx.helper.make_node("Relu", inputs=["in3"], outputs=["out3"])
@@ -1205,3 +1205,48 @@ def test_set_name_for_dimension():
     with pytest.raises(Exception) as e:
         model.set_name_for_dimension(one_const, 0, dim_name)
     assert "ONNX initializer shape dimension cannot be dynamic." in str(e)
+
+
+def test_set_input_partial_shape_using_input_edge():
+    skip_if_onnx_frontend_is_disabled()
+    fe = fem.load_by_framework(framework=ONNX_FRONTEND_NAME)
+    model = fe.load("input_model.onnx")
+
+    add_operator = model.get_place_by_operation_name("onnx_add_op")
+    add_input_edge = add_operator.get_input_port(inputPortIndex=0)
+    model.set_partial_shape(add_input_edge, PartialShape([10, 10]))
+    add_input_edge = add_operator.get_input_port(inputPortIndex=1)
+    model.set_partial_shape(add_input_edge, PartialShape([1]))
+
+    ov_model = fe.convert(model)
+    assert ov_model.input("in1").get_partial_shape() == PartialShape([10, 10])
+    assert ov_model.input("in2").get_partial_shape() == PartialShape([1])
+
+    assert ov_model.output("out4").get_partial_shape() == PartialShape([10, 10])
+
+
+def test_get_partial_shape_using_input_edge():
+    skip_if_onnx_frontend_is_disabled()
+    fe = fem.load_by_framework(framework=ONNX_FRONTEND_NAME)
+    model = fe.load("input_model.onnx")
+
+    add_operator = model.get_place_by_operation_name("onnx_add_op")
+    add_input_edge = add_operator.get_input_port(inputPortIndex=0)
+
+    pshape = model.get_partial_shape(add_input_edge)
+    assert pshape == PartialShape([2, 2])
+
+
+def test_get_partial_shape_using_output_edge():
+    skip_if_onnx_frontend_is_disabled()
+    fe = fem.load_by_framework(framework=ONNX_FRONTEND_NAME)
+    model = fe.load("input_model.onnx")
+
+    add_operator = model.get_place_by_operation_name("onnx_add_op")
+    add_output_edge = add_operator.get_output_port(outputPortIndex=0)
+
+    assert model.get_partial_shape(add_output_edge) == PartialShape([2, 2])
+
+    split_operator = model.get_place_by_tensor_name("out1").get_producing_operation()
+    out2_edge = split_operator.get_output_port(outputPortIndex=1)
+    assert model.get_partial_shape(out2_edge) == PartialShape([1, 2])

src/core/tests/onnx/onnx_editor.cpp

@@ -250,7 +250,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_head_cut) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge(1, 0)}}, {});
+    editor.extract_subgraph({{InputEdge(1, 0)}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -265,7 +265,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_head_cut_ins_and_outs) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge(1, 0)}}, {{OutputEdge(2, 0)}});
+    editor.extract_subgraph({{InputEdge(1, 0)}}, {{OutputEdge(2, 0)}});
 
     // expected to behave the same way as subgraph__linear_model_head_cut
     const auto ref_model =
@@ -281,7 +281,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_deeper_head_cut) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge(2, 0)}}, {});
+    editor.extract_subgraph({{InputEdge(2, 0)}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -296,7 +296,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_tail_cut) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({}, {{OutputEdge{1, 0}}});
+    editor.extract_subgraph({}, {{OutputEdge{1, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -311,7 +311,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_tail_cut_ins_and_outs) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{0, 0}}}, {{OutputEdge{1, 0}}});
+    editor.extract_subgraph({{InputEdge{0, 0}}}, {{OutputEdge{1, 0}}});
 
     // expected to behave the same way as subgraph__linear_model_tail_cut
     const auto ref_model =
@@ -328,7 +328,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_with_initializer_tail_cut) {
         ngraph::file_util::path_join(SERIALIZED_ZOO,
                                      "onnx/model_editor/subgraph__inception_head_with_initializer.onnx")};
 
-    editor.cut_graph_fragment({}, {{OutputEdge{1, 0}}});
+    editor.extract_subgraph({}, {{OutputEdge{1, 0}}});
 
     const auto ref_model = ngraph::file_util::path_join(
         SERIALIZED_ZOO,
@@ -344,7 +344,7 @@ NGRAPH_TEST(onnx_editor, subgraph__initializer_without_matching_input_tail_cut)
         ngraph::file_util::path_join(SERIALIZED_ZOO,
                                      "onnx/model_editor/subgraph__initializer_without_matching_input.onnx")};
 
-    editor.cut_graph_fragment({}, {{OutputEdge{1, 0}}});
+    editor.extract_subgraph({}, {{OutputEdge{1, 0}}});
 
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
                                                         "onnx/model_editor/reference/"
@@ -359,7 +359,7 @@ NGRAPH_TEST(onnx_editor, subgraph__linear_model_deeper_tail_cut) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
-    editor.cut_graph_fragment({}, {{OutputEdge{0, 0}}});
+    editor.extract_subgraph({}, {{OutputEdge{0, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -376,7 +376,7 @@ NGRAPH_TEST(onnx_editor, subgraph__no_input_params) {
 
     ONNXModelEditor editor{model_path};
 
-    editor.cut_graph_fragment({}, {});
+    editor.extract_subgraph({}, {});
 
     const auto result = compare_onnx_models(editor.model_string(), model_path);
 
@@ -388,7 +388,7 @@ NGRAPH_TEST(onnx_editor, subgraph__initializer_to_input_replacement) {
         ngraph::file_util::path_join(SERIALIZED_ZOO,
                                      "onnx/model_editor/subgraph__inception_head_with_initializer.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{0, 2}}}, {{OutputEdge{0, 0}}});
+    editor.extract_subgraph({{InputEdge{0, 2}}}, {{OutputEdge{0, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -404,7 +404,7 @@ NGRAPH_TEST(onnx_editor, subgraph__initializer_to_input_replacement_2) {
         ngraph::file_util::path_join(SERIALIZED_ZOO,
                                      "onnx/model_editor/subgraph__initializer_without_matching_input.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{0, 2}}}, {{OutputEdge{0, 0}}});
+    editor.extract_subgraph({{InputEdge{0, 2}}}, {{OutputEdge{0, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -419,7 +419,7 @@ NGRAPH_TEST(onnx_editor, subgraph__multiout_op_output_edge) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({}, {{OutputEdge{5, 1}}});
+    editor.extract_subgraph({}, {{OutputEdge{5, 1}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -434,7 +434,7 @@ NGRAPH_TEST(onnx_editor, subgraph__existing_inputs_and_outputs_based_extraction)
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{1, 1}, InputEdge{2, 0}}}, {{OutputEdge{4, 0}}});
+    editor.extract_subgraph({{InputEdge{1, 1}, InputEdge{2, 0}}}, {{OutputEdge{4, 0}}});
 
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
                                                         "onnx/model_editor/reference/"
@@ -448,7 +448,7 @@ NGRAPH_TEST(onnx_editor, subgraph__existing_inputs_and_outputs_based_extraction)
 NGRAPH_TEST(onnx_editor, subgraph__twice_input_edge_from_tensor_with_single_consumer) {
     ONNXModelEditor editor{ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/add_ab.onnx")};
 
-    editor.cut_graph_fragment({InputEdge{1, 1}}, {});
+    editor.extract_subgraph({InputEdge{1, 1}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -464,7 +464,7 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{1, 0}, InputEdge{6, 0}}}, {{OutputEdge{6, 0}, OutputEdge{4, 0}}});
+    editor.extract_subgraph({{InputEdge{1, 0}, InputEdge{6, 0}}}, {{OutputEdge{6, 0}, OutputEdge{4, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -480,7 +480,7 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{3, 0}, InputEdge{3, 1}}}, {{OutputEdge{3, 0}, OutputEdge{4, 0}}});
+    editor.extract_subgraph({{InputEdge{3, 0}, InputEdge{3, 1}}}, {{OutputEdge{3, 0}, OutputEdge{4, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -496,7 +496,7 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{3, 0}, InputEdge{6, 0}}}, {{OutputEdge{6, 0}, OutputEdge{5, 1}}});
+    editor.extract_subgraph({{InputEdge{3, 0}, InputEdge{6, 0}}}, {{OutputEdge{6, 0}, OutputEdge{5, 1}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -512,7 +512,7 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{1, 0}, InputEdge{3, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{1, 0}, InputEdge{3, 0}}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -528,7 +528,7 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({InputEdge{3, 0}}, {{OutputEdge{6, 0}, OutputEdge{5, 1}}});
+    editor.extract_subgraph({InputEdge{3, 0}}, {{OutputEdge{6, 0}, OutputEdge{5, 1}}});
 
     // expected to behave the same way as the test above
     const auto ref_model =
@@ -545,8 +545,8 @@ NGRAPH_TEST(onnx_editor, subgraph__input_edge_from_tensor_with_multiple_consumer
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{1, 0, "new_name_1"}, InputEdge{6, 0, "new_name_2"}}},
+    editor.extract_subgraph({{InputEdge{1, 0, "new_name_1"}, InputEdge{6, 0, "new_name_2"}}},
-                              {{OutputEdge{6, 0}, OutputEdge{4, 0}}});
+                            {{OutputEdge{6, 0}, OutputEdge{4, 0}}});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -562,7 +562,7 @@ NGRAPH_TEST(onnx_editor, subgraph__multiple_consumers_of_graph_input_relu2) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests_2.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{4, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{4, 0}}}, {});
 
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
                                                         "onnx/model_editor/reference/"
@@ -577,7 +577,7 @@ NGRAPH_TEST(onnx_editor, subgraph__multiple_consumers_of_graph_initializer) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests_2.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{2, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{2, 0}}}, {});
 
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
                                                         "onnx/model_editor/reference/"
@@ -592,7 +592,7 @@ NGRAPH_TEST(onnx_editor, subgraph__multiple_consumers_of_graph_initializer_2) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests_2.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{2, 0}, InputEdge{3, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{2, 0}, InputEdge{3, 0}}}, {});
 
     // same as above
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -608,7 +608,7 @@ NGRAPH_TEST(onnx_editor, subgraph__multiple_consumers_of_graph_initializer_relu2
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests_2.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{5, 0}, InputEdge{3, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{5, 0}, InputEdge{3, 0}}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -626,7 +626,7 @@ NGRAPH_TEST(onnx_editor, subgraph__invalid_edge_idx) {
 
     ONNXModelEditor editor{model_path};
     try {
-        editor.cut_graph_fragment({{InputEdge{15, 0}}}, {});
+        editor.extract_subgraph({{InputEdge{15, 0}}}, {});
     } catch (const std::exception& e) {
         std::string msg{e.what()};
         EXPECT_TRUE(msg.find("The specified node index is out of range of nodes in the original model") !=
@@ -640,7 +640,7 @@ NGRAPH_TEST(onnx_editor, subgraph__invalid_port_idx) {
 
     ONNXModelEditor editor{model_path};
     try {
-        editor.cut_graph_fragment({{InputEdge{0, 3}}}, {});
+        editor.extract_subgraph({{InputEdge{0, 3}}}, {});
     } catch (const std::exception& e) {
         std::string msg{e.what()};
         EXPECT_TRUE(msg.find("The specified node with index: 0 has not input port with index: 3") != std::string::npos);
@@ -653,7 +653,7 @@ NGRAPH_TEST(onnx_editor, subgraph__inputs_getter) {
 
     EXPECT_EQ(editor.model_inputs(), (std::vector<std::string>{"data_0", "conv1/7x7_s2_w_0", "conv1/7x7_s2_b_0"}));
 
-    editor.cut_graph_fragment({{InputEdge{1, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{1, 0}}}, {});
 
     EXPECT_EQ(editor.model_inputs(), (std::vector<std::string>{"conv1/7x7_s2_1"}));
 }
@@ -664,7 +664,7 @@ NGRAPH_TEST(onnx_editor, subgraph__custom_input_name_already_exist) {
 
     ONNXModelEditor editor{model_path};
     try {
-        editor.cut_graph_fragment({{InputEdge{1, 0, "conv1/7x7_s2_b_0"}}}, {});
+        editor.extract_subgraph({{InputEdge{1, 0, "conv1/7x7_s2_b_0"}}}, {});
     } catch (const std::exception& e) {
         std::string msg{e.what()};
         EXPECT_TRUE(msg.find("New custom input name: conv1/7x7_s2_b_0 already exist in the graph") !=
@@ -1050,7 +1050,7 @@ NGRAPH_TEST(onnx_editor, editor_api_use_edge_mapper_with_graph_cutter) {
 
     const auto output_edge = editor.find_output_edge(EditorNode(EditorOutput("mul2")), EditorOutput(0));
     // OutputEdge{4, "mul2"}
-    editor.cut_graph_fragment({input_edge_1, input_edge_2}, {output_edge});
+    editor.extract_subgraph({input_edge_1, input_edge_2}, {output_edge});
 
     const auto ref_model = ngraph::file_util::path_join(SERIALIZED_ZOO,
                                                         "onnx/model_editor/reference/"
@@ -1082,7 +1082,7 @@ NGRAPH_TEST(onnx_editor, editor_api_use_edge_mapper_with_graph_cutter_custom_nam
     const auto input_edge_2 =
         editor.find_input_edge(EditorNode{EditorOutput{"split2"}}, EditorInput{"add2", "new_name_2"});
 
-    editor.cut_graph_fragment({input_edge_1, input_edge_2}, {});
+    editor.extract_subgraph({input_edge_1, input_edge_2}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -1201,7 +1201,7 @@ NGRAPH_TEST(onnx_editor, editor_api_input_edge_from_tensor_with_single_consumer)
     ONNXModelEditor editor{ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/add_ab.onnx")};
 
     const auto edge = editor.find_input_edge(EditorNode{EditorOutput{"Y"}}, EditorInput{1});
-    editor.cut_graph_fragment({edge}, {});
+    editor.extract_subgraph({edge}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO,
@@ -1346,7 +1346,7 @@ NGRAPH_TEST(onnx_editor, combined__cut_and_replace_shape) {
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph__inception_head.onnx")};
 
     const auto new_shape = PartialShape({1, 64, 112, 112});
-    editor.cut_graph_fragment({{InputEdge(1, 0)}}, {});
+    editor.extract_subgraph({{InputEdge(1, 0)}}, {});
     editor.set_input_shapes({{"conv1/7x7_s2_1", new_shape}});
 
     const auto ref_model =
@@ -1365,7 +1365,7 @@ NGRAPH_TEST(onnx_editor, cut_operator_with_no_schema) {
     ONNXModelEditor editor{
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/unknown_input_value_info.onnx")};
 
-    editor.cut_graph_fragment({{InputEdge{1, 0}}}, {});
+    editor.extract_subgraph({{InputEdge{1, 0}}}, {});
 
     const auto ref_model =
         ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/reference/unknown_input_value_info.onnx");

src/frontends/onnx/frontend/src/editor.cpp

@@ -336,8 +336,8 @@ PartialShape onnx_editor::ONNXModelEditor::get_tensor_shape(const std::string& t
     }
 }
 
-void onnx_editor::ONNXModelEditor::cut_graph_fragment(const std::vector<InputEdge>& inputs,
+void onnx_editor::ONNXModelEditor::extract_subgraph(const std::vector<InputEdge>& inputs,
-                                                      const std::vector<OutputEdge>& outputs) {
+                                                    const std::vector<OutputEdge>& outputs) {
     if (inputs.empty() && outputs.empty()) {
         return;
     }
@@ -494,6 +494,16 @@ void onnx_editor::ONNXModelEditor::set_node_name(const EditorNode& node, const s
     *graph->mutable_node(node_idx)->mutable_name() = new_name;
 }
 
+std::string onnx_editor::ONNXModelEditor::get_node_name(const EditorNode& node) const {
+    if (node.m_node_index >= 0) {
+        OPENVINO_ASSERT(node.m_node_index < m_pimpl->m_model_proto->graph().node().size(), "XXX");
+        const auto& n = m_pimpl->m_model_proto->graph().node(node.m_node_index);
+        return n.has_name() ? n.name() : "";
+    } else {
+        return node.m_node_name;
+    }
+}
+
 void onnx_editor::ONNXModelEditor::clear_nodes_name(const std::string& name) {
     const auto graph = m_pimpl->m_model_proto->mutable_graph();
 

src/frontends/onnx/frontend/src/editor.hpp

@@ -80,7 +80,7 @@ public:
     ///
     /// \param inputs A collection of input edges which become new inputs to the graph
     /// \param outputs A collection of output edges which become new outputs of the graph
-    void cut_graph_fragment(const std::vector<InputEdge>& inputs, const std::vector<OutputEdge>& outputs);
+    void extract_subgraph(const std::vector<InputEdge>& inputs, const std::vector<OutputEdge>& outputs);
 
     /// \brief Modifies the in-memory representation of the model by setting custom input
     ///        values for inputs specified in the provided map.
@@ -112,6 +112,11 @@ public:
     /// \param new_name New name of the node.
     void set_node_name(const EditorNode& node, const std::string& new_name);
 
+    /// \brief Retrieves a node name from the in-memory ONNX model.
+    ///
+    /// \param node Node descriptor for which the lookup should be performed.
+    std::string get_node_name(const EditorNode& node) const;
+
     /// \brief Removes node name for all nodes with given name.
     ///
     /// \note Empty and not present names are accepted.

src/frontends/onnx/frontend/src/input_model.cpp

@@ -123,18 +123,46 @@ void InputModel::free_name_for_tensor(const std::string&) {
 }
 
 void InputModel::set_partial_shape(const ov::frontend::Place::Ptr& place, const ngraph::PartialShape& shape) {
-    std::map<std::string, ngraph::PartialShape> m;
+    std::string input_name;  // name of the model input which should be reshaped
-    m[place->get_names()[0]] = shape;
+    const auto input_edge = std::dynamic_pointer_cast<PlaceInputEdge>(place);
-    m_editor->set_input_shapes(m);
+    if (input_edge) {
+        const auto tensor_names = input_edge->get_source_tensor()->get_names();
+        OPENVINO_ASSERT(!tensor_names.empty(), "Cannot retrieve input name. Setting new input shape is not possible.");
+        input_name = tensor_names[0];
+    } else {
+        // fallback in case something else than an InputEdge is passed in - try to retrieve its name and reshape
+        OPENVINO_ASSERT(!place->get_names().empty(),
+                        "Cannot retrieve input name. Setting new input shape is not possible.");
+        input_name = place->get_names()[0];
+    }
+
+    m_editor->set_input_shapes({{input_name, shape}});
 }
 
 ngraph::PartialShape InputModel::get_partial_shape(const ov::frontend::Place::Ptr& place) const {
-    return m_editor->get_tensor_shape(place->get_names().at(0));
+    std::string tensor_name;  // name of the model input which should be reshaped
+    const auto input_edge = std::dynamic_pointer_cast<PlaceInputEdge>(place);
+    const auto output_edge = std::dynamic_pointer_cast<PlaceOutputEdge>(place);
+    if (input_edge) {
+        const auto tensor_names = input_edge->get_source_tensor()->get_names();
+        OPENVINO_ASSERT(!tensor_names.empty(),
+                        "Cannot retrieve source tensor name for this InputEdge and thus partial shape.");
+        tensor_name = tensor_names[0];
+    } else if (output_edge) {
+        const auto tensor_names = output_edge->get_target_tensor()->get_names();
+        OPENVINO_ASSERT(!tensor_names.empty(),
+                        "Cannot retrieve target tensor name for this OutputEdge and thus partial shape.");
+        tensor_name = tensor_names[0];
+    } else {
+        tensor_name = place->get_names().at(0);
+    }
+
+    return m_editor->get_tensor_shape(tensor_name);
 }
 
 void InputModel::set_element_type(const ov::frontend::Place::Ptr& place, const ngraph::element::Type& type) {
     std::map<std::string, ngraph::element::Type_t> m;
-    m[place->get_names()[0]] = type;
+    m[place->get_names().at(0)] = type;
     m_editor->set_input_types(m);
 }
 
@@ -179,7 +207,7 @@ void InputModel::extract_subgraph(const std::vector<ov::frontend::Place::Ptr>& i
         if (const auto input_port = std::dynamic_pointer_cast<PlaceInputEdge>(input)) {
             onnx_inputs.push_back(input_port->get_input_edge());
         } else if (const auto tensor = std::dynamic_pointer_cast<PlaceTensor>(input)) {
-            auto name = tensor->get_names()[0];
+            const auto name = tensor->get_names().at(0);
             const auto consumers = m_editor->find_output_consumers(name);
             std::transform(std::begin(consumers),
                            std::end(consumers),
@@ -224,5 +252,5 @@ void InputModel::extract_subgraph(const std::vector<ov::frontend::Place::Ptr>& i
                            });
         }
     }
-    m_editor->cut_graph_fragment(onnx_inputs, onnx_outputs);
+    m_editor->extract_subgraph(onnx_inputs, onnx_outputs);
 }

src/frontends/onnx/frontend/src/place.cpp

@@ -21,6 +21,10 @@ onnx_editor::InputEdge PlaceInputEdge::get_input_edge() const {
     return m_edge;
 }
 
+std::vector<std::string> PlaceInputEdge::get_names() const {
+    return {"InputEdge{" + std::to_string(m_edge.m_node_idx) + ", " + std::to_string(m_edge.m_port_idx) + "}"};
+}
+
 bool PlaceInputEdge::is_input() const {
     return m_editor->is_input(m_edge);
 }
@@ -70,6 +74,10 @@ onnx_editor::OutputEdge PlaceOutputEdge::get_output_edge() const {
     return m_edge;
 }
 
+std::vector<std::string> PlaceOutputEdge::get_names() const {
+    return {"OutputEdge{" + std::to_string(m_edge.m_node_idx) + ", " + std::to_string(m_edge.m_port_idx) + "}"};
+}
+
 bool PlaceOutputEdge::is_input() const {
     return false;
 }
@@ -198,7 +206,11 @@ PlaceOp::PlaceOp(onnx_editor::EditorNode&& node, std::shared_ptr<onnx_editor::ON
       m_editor{std::move(editor)} {}
 
 std::vector<std::string> PlaceOp::get_names() const {
-    return {m_node.m_node_name};
+    if (!m_node.m_node_name.empty()) {
+        return {m_node.m_node_name};
+    } else {
+        return {m_editor->get_node_name(m_node)};
+    }
 }
 
 const onnx_editor::EditorNode& PlaceOp::get_editor_node() const {

src/frontends/onnx/frontend/src/place.hpp

@@ -22,6 +22,7 @@ public:
     onnx_editor::InputEdge get_input_edge() const;
 
     // external usage
+    std::vector<std::string> get_names() const override;
     bool is_input() const override;
     bool is_output() const override;
     bool is_equal(const Place::Ptr& another) const override;
@@ -45,6 +46,7 @@ public:
     onnx_editor::OutputEdge get_output_edge() const;
 
     // external usage
+    std::vector<std::string> get_names() const override;
     bool is_input() const override;
     bool is_output() const override;
     bool is_equal(const Place::Ptr& another) const override;

tools/mo/openvino/tools/mo/moc_frontend/extractor.py

@@ -32,8 +32,7 @@ def decode_name_with_port(input_model: InputModel, node_name: str):
     if match_post:
         node_post = input_model.get_place_by_operation_name(match_post.group(1))
         if node_post:
-            node_post = node_post.get_output_port(
+            node_post = node_post.get_output_port(outputPortIndex=int(match_post.group(2)))
-                outputPortIndex=int(match_post.group(2)))
             if node_post:
                 found_node_names.append(match_post.group(1))
                 found_nodes.append(node_post)
@@ -43,8 +42,7 @@ def decode_name_with_port(input_model: InputModel, node_name: str):
     if match_pre:
         node_pre = input_model.get_place_by_operation_name(match_pre.group(2))
         if node_pre:
-            node_pre = node_pre.get_input_port(
+            node_pre = node_pre.get_input_port(inputPortIndex=int(match_pre.group(1)))
-                inputPortIndex=int(match_pre.group(1)))
             if node_pre:
                 found_node_names.append(match_pre.group(2))
                 found_nodes.append(node_pre)

tools/mo/openvino/tools/mo/moc_frontend/pipeline.py

@@ -55,10 +55,20 @@ def moc_pipeline(argv: argparse.Namespace, moc_front_end: FrontEnd):
         new_output_places = [x['node'] for x in outputs]
         log.debug('Using extract subgraph')
         input_model.extract_subgraph(new_input_places, new_output_places)
+        # invalidation of existing Place objects could have happened in the operation above
+        if user_shapes:
+            user_shapes, outputs, freeze_placeholder = fe_user_data_repack(
+                input_model, argv.placeholder_shapes, argv.placeholder_data_types,
+                argv.output, argv.freeze_placeholder_with_value)
     elif not inputs_equal:
         new_input_places = [x['node'] for x in user_shapes]
         log.debug('Using override_all_inputs')
         input_model.override_all_inputs(new_input_places)
+        # invalidation of existing Place objects could have happened in the operation above
+        if user_shapes:
+            user_shapes, outputs, freeze_placeholder = fe_user_data_repack(
+                input_model, argv.placeholder_shapes, argv.placeholder_data_types,
+                argv.output, argv.freeze_placeholder_with_value)
     elif not outputs_equal:
         new_output_places = [x['node'] for x in outputs]
         log.debug('Using override_all_outputs')