diff --git a/src/bindings/python/src/pyopenvino/frontend/onnx/extension.cpp b/src/bindings/python/src/pyopenvino/frontend/onnx/extension.cpp index 736085c220a..a3926441aa6 100644 --- a/src/bindings/python/src/pyopenvino/frontend/onnx/extension.cpp +++ b/src/bindings/python/src/pyopenvino/frontend/onnx/extension.cpp @@ -28,10 +28,12 @@ void regclass_frontend_onnx_ConversionExtension(py::module m) { public: using Ptr = std::shared_ptr; using PyCreatorFunction = std::function; - PyConversionExtension(const std::string& op_type, const PyCreatorFunction& f) - : ConversionExtension(op_type, [f](const ov::frontend::NodeContext& node) -> ov::OutputVector { + PyConversionExtension(const std::string& op_type, const std::string& op_domain, const PyCreatorFunction& f) + : ConversionExtension(op_type, op_domain, [f](const ov::frontend::NodeContext& node) -> ov::OutputVector { return f(static_cast(&node)); }) {} + PyConversionExtension(const std::string& op_type, const PyCreatorFunction& f) + : PyConversionExtension(op_type, "", f) {} }; py::class_ ext(m, "ConversionExtensionONNX", @@ -41,6 +43,10 @@ void regclass_frontend_onnx_ConversionExtension(py::module m) { return std::make_shared(op_type, f); })); + ext.def(py::init([](const std::string& op_type, const std::string& op_domain, const PyConversionExtension::PyCreatorFunction& f) { + return std::make_shared(op_type, op_domain, f); + })); + ext.def_property_readonly_static("m_converter", &ConversionExtension::get_converter); } @@ -78,4 +84,22 @@ void regclass_frontend_onnx_OpExtension(py::module m) { py::arg("fw_type_name"), py::arg("attr_names_map") = std::map(), py::arg("attr_values_map") = std::map()); + + ext.def(py::init([](const std::string& ov_type_name, + const std::string& fw_type_name, + const std::string& fw_domain, + const std::map& attr_names_map, + const std::map& attr_values_map) { + + std::map any_map; + for (const auto& it : attr_values_map) { + any_map[it.first] = py_object_to_any(it.second); + } + return std::make_shared>(ov_type_name, fw_type_name, fw_domain, attr_names_map, any_map); + }), + py::arg("ov_type_name"), + py::arg("fw_type_name"), + py::arg("fw_domain"), + py::arg("attr_names_map") = std::map(), + py::arg("attr_values_map") = std::map()); } diff --git a/src/bindings/python/tests/test_frontend/test_frontend_onnx.py b/src/bindings/python/tests/test_frontend/test_frontend_onnx.py index c945b07a15c..22a8b72f578 100644 --- a/src/bindings/python/tests/test_frontend/test_frontend_onnx.py +++ b/src/bindings/python/tests/test_frontend/test_frontend_onnx.py @@ -142,6 +142,24 @@ def create_onnx_model_for_op_extension(): return make_model(graph, producer_name="ONNX Frontend") +def create_onnx_model_extension_with_custom_domain(): + add = onnx.helper.make_node("CustomAdd", inputs=["x", "y"], outputs=["z"], domain="custom_domain") + const_tensor = onnx.helper.make_tensor("const_tensor", + onnx.TensorProto.FLOAT, + (2, 2), + [0.5, 1, 1.5, 2.0]) + const_node = onnx.helper.make_node("Constant", [], outputs=["const_node"], + value=const_tensor, name="const_node") + mul = onnx.helper.make_node("Mul", inputs=["z", "const_node"], outputs=["out"]) + input_tensors = [ + make_tensor_value_info("x", onnx.TensorProto.FLOAT, (2, 2)), + make_tensor_value_info("y", onnx.TensorProto.FLOAT, (2, 2)), + ] + output_tensors = [make_tensor_value_info("out", onnx.TensorProto.FLOAT, (2, 2))] + graph = make_graph([add, const_node, mul], "graph", input_tensors, output_tensors) + return make_model(graph, producer_name="ONNX Frontend") + + def run_model(model, *inputs, expected): runtime = get_runtime() computation = runtime.computation(model) @@ -159,6 +177,7 @@ onnx_model_2_filename = "model2.onnx" onnx_model_with_custom_attributes_filename = "model_custom_attributes.onnx" onnx_model_with_subgraphs_filename = "model_subgraphs.onnx" onnx_model_for_op_extension_test = "model_op_extension.onnx" +onnx_model_extension_with_custom_domain = "model_extension_custom_domain.onnx" ONNX_FRONTEND_NAME = "onnx" @@ -169,6 +188,7 @@ def setup_module(): onnx_model_with_custom_attributes_filename) onnx.save_model(create_onnx_model_with_subgraphs(), onnx_model_with_subgraphs_filename) onnx.save_model(create_onnx_model_for_op_extension(), onnx_model_for_op_extension_test) + onnx.save_model(create_onnx_model_extension_with_custom_domain(), onnx_model_extension_with_custom_domain) def teardown_module(): @@ -177,6 +197,7 @@ def teardown_module(): os.remove(onnx_model_with_custom_attributes_filename) os.remove(onnx_model_with_subgraphs_filename) os.remove(onnx_model_for_op_extension_test) + os.remove(onnx_model_extension_with_custom_domain) def skip_if_onnx_frontend_is_disabled(): @@ -482,6 +503,53 @@ def test_onnx_conversion_extension(): assert invoked +def test_onnx_conversion_extension_with_custom_domain(): + skip_if_onnx_frontend_is_disabled() + + # use specific (openvino.frontend.onnx) import here + from openvino.frontend.onnx import ConversionExtension + from openvino.frontend import NodeContext + import openvino.runtime.opset8 as ops + + fe = fem.load_by_model(onnx_model_extension_with_custom_domain) + assert fe + assert fe.get_name() == "onnx" + + invoked = False + + def custom_converter(node: NodeContext): + nonlocal invoked + invoked = True + input_1 = node.get_input(0) + input_2 = node.get_input(1) + add = ops.add(input_1, input_2) + return [add.output(0)] + + fe.add_extension(ConversionExtension("CustomAdd", "custom_domain", custom_converter)) + input_model = fe.load(onnx_model_extension_with_custom_domain) + assert input_model + model = fe.convert(input_model) + assert model + assert invoked + + +def test_onnx_op_extension_with_custom_domain(): + skip_if_onnx_frontend_is_disabled() + + # use specific (openvino.frontend.onnx) import here + from openvino.frontend.onnx import OpExtension + + fe = fem.load_by_model(onnx_model_extension_with_custom_domain) + assert fe + assert fe.get_name() == "onnx" + + fe.add_extension(OpExtension("opset1.Add", "CustomAdd", "custom_domain", {}, {"auto_broadcast": "numpy"})) + input_model = fe.load(onnx_model_extension_with_custom_domain) + assert input_model + model = fe.convert(input_model) + assert model + + @pytest.mark.parametrize("opset_prefix", ["opset1.", "opset1::", "opset8.", "opset8::", ""]) def test_op_extension_specify_opset(opset_prefix): skip_if_onnx_frontend_is_disabled() diff --git a/src/frontends/common/include/openvino/frontend/extension/op.hpp b/src/frontends/common/include/openvino/frontend/extension/op.hpp index ef60c29e36b..13da3d4e559 100644 --- a/src/frontends/common/include/openvino/frontend/extension/op.hpp +++ b/src/frontends/common/include/openvino/frontend/extension/op.hpp @@ -31,6 +31,72 @@ inline const ov::OpSet& get_opset_by_name(const std::string& opset_name) { } } +/// \brief The helper function to create an instance of ov::Node class initialized by provided type name. +/// Expected formats: +/// - opsetN::OpName +/// - opsetN.OpName +/// - OpName +/// \param ov_type_name Type name of created ov::Node. +inline std::shared_ptr create_ov_node_by_name(const std::string& ov_type_name) { + auto split = [](const std::string& s, const std::string& delimiter) { + size_t pos_start = 0, pos_end, delim_len = delimiter.length(); + std::string token; + std::vector res; + + while ((pos_end = s.find(delimiter, pos_start)) != std::string::npos) { + token = s.substr(pos_start, pos_end - pos_start); + pos_start = pos_end + delim_len; + res.push_back(token); + } + + res.push_back(s.substr(pos_start)); + return res; + }; + + std::string opset_name; + std::string op_name; + auto cnt_colons = std::count(ov_type_name.begin(), ov_type_name.end(), ':'); + auto cnt_dots = std::count(ov_type_name.begin(), ov_type_name.end(), '.'); + if (cnt_colons == 2 && cnt_dots == 0) { + auto divided = split(ov_type_name, "::"); + if (divided.size() != 2) { + FRONT_END_GENERAL_CHECK(false, + "Invalid OpenVINO operation format, one of the next is expected:" + "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", + ov_type_name); + } + opset_name = divided[0]; + op_name = divided[1]; + } else if (cnt_colons == 0 && cnt_dots == 1) { + auto divided = split(ov_type_name, "."); + if (divided.size() != 2) { + FRONT_END_GENERAL_CHECK(false, + "Invalid OpenVINO operation format, one of the next is expected:" + "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", + ov_type_name); + } + opset_name = divided[0]; + op_name = divided[1]; + } else if (cnt_colons == 0 && cnt_dots == 0) { + opset_name = "latest"; + op_name = ov_type_name; + } else { + FRONT_END_GENERAL_CHECK(false, + "Invalid OpenVINO operation format, one of the next is expected: \n" + "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", + ov_type_name); + } + + const auto& opset = get_opset_by_name(opset_name); + if (!opset.contains_type(op_name)) { + FRONT_END_GENERAL_CHECK(false, + "OpenVINO opset doesn't contain operation with " + "name ", + op_name); + } + return std::shared_ptr(opset.create(op_name)); +} + // One-to-one operation mapping for OVOpType != void which means OV type is specified by OVOpType // See a specialization for OVOptype = void template @@ -139,72 +205,8 @@ OpExtensionBase::OpExtensionBase(const std::string& ov const std::map& attr_values_map) : BaseConversionType(fw_type_name, OpConversionFunction( - [=]() -> std::shared_ptr { - auto split = [](const std::string& s, const std::string& delimiter) { - size_t pos_start = 0, pos_end, delim_len = delimiter.length(); - std::string token; - std::vector res; - - while ((pos_end = s.find(delimiter, pos_start)) != std::string::npos) { - token = s.substr(pos_start, pos_end - pos_start); - pos_start = pos_end + delim_len; - res.push_back(token); - } - - res.push_back(s.substr(pos_start)); - return res; - }; - - // Expected formats: - // opsetN::OpName - // opsetN.OpName - // OpName - std::string opset_name; - std::string op_name; - auto cnt_colons = std::count(ov_type_name.begin(), ov_type_name.end(), ':'); - auto cnt_dots = std::count(ov_type_name.begin(), ov_type_name.end(), '.'); - if (cnt_colons == 2 && cnt_dots == 0) { - auto divided = split(ov_type_name, "::"); - if (divided.size() != 2) { - FRONT_END_GENERAL_CHECK( - false, - "Invalid OpenVINO operation format, one of the next is expected:" - "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", - ov_type_name); - } - opset_name = divided[0]; - op_name = divided[1]; - } else if (cnt_colons == 0 && cnt_dots == 1) { - auto divided = split(ov_type_name, "."); - if (divided.size() != 2) { - FRONT_END_GENERAL_CHECK( - false, - "Invalid OpenVINO operation format, one of the next is expected:" - "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", - ov_type_name); - } - opset_name = divided[0]; - op_name = divided[1]; - } else if (cnt_colons == 0 && cnt_dots == 0) { - opset_name = "latest"; - op_name = ov_type_name; - } else { - FRONT_END_GENERAL_CHECK( - false, - "Invalid OpenVINO operation format, one of the next is expected: \n" - "opsetN::OpName or opsetN.OpName or OpName. Provided operation format: ", - ov_type_name); - } - - const auto& opset = get_opset_by_name(opset_name); - if (!opset.contains_type(op_name)) { - FRONT_END_GENERAL_CHECK(false, - "OpenVINO opset doesn't contain operation with " - "name ", - op_name); - } - - return std::shared_ptr(opset.create(op_name)); + [ov_type_name]() { + return create_ov_node_by_name(ov_type_name); }, attr_names_map, attr_values_map)) {} diff --git a/src/frontends/onnx/frontend/include/openvino/frontend/onnx/extension/op.hpp b/src/frontends/onnx/frontend/include/openvino/frontend/onnx/extension/op.hpp index 6285b435ab3..0d9d97bca08 100644 --- a/src/frontends/onnx/frontend/include/openvino/frontend/onnx/extension/op.hpp +++ b/src/frontends/onnx/frontend/include/openvino/frontend/onnx/extension/op.hpp @@ -11,7 +11,61 @@ namespace frontend { namespace onnx { template -using OpExtension = ov::frontend::OpExtensionBase; +class OpExtension : public ConversionExtension { +public: + OpExtension(const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : OpExtension(OVOpType::get_type_info_static().name, "", attr_names_map, attr_values_map) {} + + OpExtension(const std::string& fw_type_name, + const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : OpExtension(fw_type_name, "", attr_names_map, attr_values_map) {} + + OpExtension(const std::string& fw_type_name, + const std::string& fw_domain, + const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : ConversionExtension(fw_type_name, + fw_domain, + OpConversionFunction( + []() { + return std::make_shared(); + }, + attr_names_map, + attr_values_map)) {} +}; + +template <> +class OpExtension : public ConversionExtension { +public: + OpExtension() = delete; + + explicit OpExtension(const std::string& fw_ov_type_name, + const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : OpExtension(fw_ov_type_name, fw_ov_type_name, attr_names_map, attr_values_map) {} + + OpExtension(const std::string& ov_type_name, + const std::string& fw_type_name, + const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : OpExtension(ov_type_name, fw_type_name, "", attr_names_map, attr_values_map) {} + + OpExtension(const std::string& ov_type_name, + const std::string& fw_type_name, + const std::string& fw_domain_name, + const std::map& attr_names_map = {}, + const std::map& attr_values_map = {}) + : ConversionExtension(fw_type_name, + fw_domain_name, + OpConversionFunction( + [ov_type_name]() { + return create_ov_node_by_name(ov_type_name); + }, + attr_names_map, + attr_values_map)) {} +}; } // namespace onnx } // namespace frontend diff --git a/src/frontends/onnx/tests/models/relu_custom_domain.prototxt b/src/frontends/onnx/tests/models/relu_custom_domain.prototxt new file mode 100644 index 00000000000..dc51e829930 --- /dev/null +++ b/src/frontends/onnx/tests/models/relu_custom_domain.prototxt @@ -0,0 +1,45 @@ +ir_version: 4 +producer_name: "OV ONNX FE" +graph { + node { + input: "A" + output: "B" + op_type: "CustomRelu" + domain: "my_custom_domain" + } + input { + name: "A" + type { + tensor_type { + elem_type: 1 + shape { + dim { + dim_value: 2 + } + dim { + dim_value: 2 + } + } + } + } + } + output { + name: "B" + type { + tensor_type { + elem_type: 1 + shape { + dim { + dim_value: 2 + } + dim { + dim_value: 2 + } + } + } + } + } +} +opset_import { + version: 1 +} diff --git a/src/frontends/onnx/tests/op_extension.cpp b/src/frontends/onnx/tests/op_extension.cpp index 67b49810900..a73d5877444 100644 --- a/src/frontends/onnx/tests/op_extension.cpp +++ b/src/frontends/onnx/tests/op_extension.cpp @@ -4,10 +4,12 @@ #include "op_extension.hpp" +#include "common_test_utils/file_utils.hpp" #include "onnx_utils.hpp" #include "openvino/frontend/extension/op.hpp" #include "openvino/frontend/onnx/extension/op.hpp" #include "openvino/frontend/onnx/frontend.hpp" +#include "openvino/op/relu.hpp" #include "so_extension.hpp" using namespace ov::frontend; @@ -141,4 +143,30 @@ INSTANTIATE_TEST_SUITE_P(ONNXOpExtensionViaONNXConstructor, INSTANTIATE_TEST_SUITE_P(ONNXOpExtensionViaCommonConstructor, FrontEndOpExtensionTest, ::testing::Values(getTestDataOpExtensionViaCommonConstructor()), - FrontEndOpExtensionTest::getTestCaseName); \ No newline at end of file + FrontEndOpExtensionTest::getTestCaseName); + +TEST(ONNXOpExtensionViaCommonConstructor, onnx_op_extension_via_template_arg_with_custom_domain) { + const auto ext = std::make_shared>("CustomRelu", "my_custom_domain"); + + auto fe = std::make_shared(); + fe->add_extension(ext); + + const auto input_model = fe->load(CommonTestUtils::getModelFromTestModelZoo( + ov::util::path_join({TEST_ONNX_MODELS_DIRNAME, "relu_custom_domain.onnx"}))); + + std::shared_ptr model; + EXPECT_NO_THROW(fe->convert(input_model)); +} + +TEST(ONNXOpExtensionViaCommonConstructor, onnx_op_extension_via_ov_type_name_with_custom_domain) { + const auto ext = std::make_shared>("opset1::Relu", "CustomRelu", "my_custom_domain"); + + auto fe = std::make_shared(); + fe->add_extension(ext); + + const auto input_model = fe->load(CommonTestUtils::getModelFromTestModelZoo( + ov::util::path_join({TEST_ONNX_MODELS_DIRNAME, "relu_custom_domain.onnx"}))); + + std::shared_ptr model; + EXPECT_NO_THROW(fe->convert(input_model)); +}