diff --git a/src/frontends/pytorch/src/frontend.cpp b/src/frontends/pytorch/src/frontend.cpp index 4d8305dde39..054accca91b 100644 --- a/src/frontends/pytorch/src/frontend.cpp +++ b/src/frontends/pytorch/src/frontend.cpp @@ -23,6 +23,7 @@ #include "transforms/aten_index_put_replacer.hpp" #include "transforms/aten_index_replacer.hpp" #include "transforms/aten_stack_list_construct_replacer.hpp" +#include "transforms/dict_resolver.hpp" #include "transforms/einsum_list_construct.hpp" #include "transforms/listconstruct_replacer.hpp" #include "transforms/min_max_prim_list_construct_replacer.hpp" @@ -112,6 +113,7 @@ void FrontEnd::normalize(const std::shared_ptr& model) const { manager.register_pass(); manager.register_pass(); manager.register_pass(); + manager.register_pass(); manager.register_pass(); manager.register_pass(); diff --git a/src/frontends/pytorch/src/transforms/dict_resolver.cpp b/src/frontends/pytorch/src/transforms/dict_resolver.cpp new file mode 100644 index 00000000000..6f1d4c4e03d --- /dev/null +++ b/src/frontends/pytorch/src/transforms/dict_resolver.cpp @@ -0,0 +1,71 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "dict_resolver.hpp" + +#include "openvino/core/rt_info.hpp" +#include "openvino/op/result.hpp" +#include "openvino/op/util/framework_node.hpp" +#include "openvino/pass/pattern/op/wrap_type.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace pass { + +using namespace ov::pass; +using namespace ov::op; + +bool DictResolver::run_on_model(const std::shared_ptr& model) { + bool changed = false; + const auto results = model->get_results(); + for (const auto& res : results) { + if (auto dict_construct_node = cast_fw_node(res->get_input_node_shared_ptr(0), "prim::DictConstruct")) { + const auto inputs = dict_construct_node->input_values(); + if (inputs.size() % 2) { + // inputs must be divisible by 2 + return false; + } + ResultVector new_outputs; + for (size_t i = 0; i < inputs.size(); i += 2) { + auto new_output = inputs.at(i + 1); + const auto name_node = inputs.at(i); + auto fw_node = std::dynamic_pointer_cast(name_node.get_node_shared_ptr()); + if (!fw_node) { + // odd inputs must contain constant strings encoded as fw nodes + return false; + } + const auto attrs = fw_node->get_attrs(); + if (attrs.find("string_value") == attrs.end()) { + // fw node must contain string value + return false; + } + auto name = attrs.at("string_value"); + new_output.add_names({name}); + new_outputs.push_back(std::make_shared(new_output)); + } + bool after_res = false; + for (const auto& result : results) { + if (after_res) { + // To preserve output order remove results after this one and insert them after new outputs + model->remove_result(result); + new_outputs.push_back(result); + } + if (result == res) { + after_res = true; + model->remove_result(result); + } + } + model->add_results(new_outputs); + changed = true; + } + } + return changed; +}; + +} // namespace pass +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/transforms/dict_resolver.hpp b/src/frontends/pytorch/src/transforms/dict_resolver.hpp new file mode 100644 index 00000000000..7cdec639cf2 --- /dev/null +++ b/src/frontends/pytorch/src/transforms/dict_resolver.hpp @@ -0,0 +1,24 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include "openvino/pass/graph_rewrite.hpp" +#include "openvino/pass/pass.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace pass { + +class DictResolver : public ov::pass::ModelPass { +public: + OPENVINO_RTTI("ov::frontend::pytorch::pass::DictResolver"); + bool run_on_model(const std::shared_ptr& model) override; +}; + +} // namespace pass +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py index a5ad7748f91..77d80af650c 100644 --- a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py +++ b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py @@ -81,7 +81,11 @@ class PytorchLayerTest: flatten_fw_res = [] - def flattenize_list_outputs(res): + def flattenize_dict_outputs(res): + if isinstance(res, dict): + return flattenize_outputs(res.values()) + + def flattenize_outputs(res): results = [] for res_item in res: # if None is at output we skip it @@ -89,13 +93,17 @@ class PytorchLayerTest: continue # If input is list or tuple flattenize it if isinstance(res_item, (list, tuple)): - decomposed_res = flattenize_list_outputs(res_item) + decomposed_res = flattenize_outputs(res_item) + results.extend(decomposed_res) + continue + if isinstance(res_item, dict): + decomposed_res = flattenize_dict_outputs(res_item) results.extend(decomposed_res) continue results.append(res_item) return results - flatten_fw_res = flattenize_list_outputs(fw_res) + flatten_fw_res = flattenize_outputs(fw_res) assert len(flatten_fw_res) == len( output_list), f'number of outputs not equal, {len(flatten_fw_res)} != {len(output_list)}' diff --git a/tests/layer_tests/pytorch_tests/test_dict.py b/tests/layer_tests/pytorch_tests/test_dict.py new file mode 100644 index 00000000000..a3e3c29b6ba --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_dict.py @@ -0,0 +1,26 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import torch + +from pytorch_layer_test_class import PytorchLayerTest + + +class TestDict(PytorchLayerTest): + + def _prepare_input(self): + return (np.random.randn(2, 5, 3, 4).astype(np.float32),) + + def create_model(self): + class aten_dict(torch.nn.Module): + def forward(self, x): + return {"b": x, "a": x + x, "c": 2 * x}, x / 2 + + return aten_dict(), None, "prim::DictConstruct" + + @pytest.mark.nightly + @pytest.mark.precommit + def test_dict(self, ie_device, precision, ir_version): + self._test(*self.create_model(), ie_device, precision, ir_version)