diff --git a/tools/mo/automation/package_BOM.txt b/tools/mo/automation/package_BOM.txt index cacf4e94123..4a0f65edfce 100644 --- a/tools/mo/automation/package_BOM.txt +++ b/tools/mo/automation/package_BOM.txt @@ -1076,6 +1076,8 @@ openvino/tools/mo/utils/ir_reader/extenders/strided_slice_extender.py openvino/tools/mo/utils/ir_reader/extenders/tensoriterator_extender.py openvino/tools/mo/utils/ir_reader/extenders/topk_extender.py openvino/tools/mo/utils/ir_reader/extenders/variadic_split_extender.py +openvino/tools/mo/utils/ir_reader/internal_ops/squeeze.py +openvino/tools/mo/utils/ir_reader/internal_ops/unsqueeze.py openvino/tools/mo/utils/ir_reader/layer_to_class.py openvino/tools/mo/utils/ir_reader/restore_graph.py openvino/tools/mo/utils/json_schema.py diff --git a/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/squeeze.py b/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/squeeze.py new file mode 100644 index 00000000000..d31b145f63f --- /dev/null +++ b/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/squeeze.py @@ -0,0 +1,14 @@ +# Copyright (C) 2018-2022 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from openvino.tools.mo.graph.graph import Node +from openvino.tools.mo.ops.squeeze import Squeeze + + +class SqueezeInternal(Squeeze): + @staticmethod + def infer(node: Node): + axis_value = node.in_port(1).data.get_value() + Squeeze.infer(node) + # preserve initial axis value + node.in_port(1).data.set_value(axis_value) diff --git a/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/unsqueeze.py b/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/unsqueeze.py new file mode 100644 index 00000000000..25afb17c4e8 --- /dev/null +++ b/tools/mo/openvino/tools/mo/utils/ir_reader/internal_ops/unsqueeze.py @@ -0,0 +1,14 @@ +# Copyright (C) 2018-2022 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from openvino.tools.mo.graph.graph import Node +from openvino.tools.mo.ops.unsqueeze import Unsqueeze + + +class UnsqueezeInternal(Unsqueeze): + @staticmethod + def infer(node: Node): + axis_value = node.in_port(1).data.get_value() + Unsqueeze.infer(node) + # preserve initial axis value + node.in_port(1).data.set_value(axis_value) diff --git a/tools/mo/openvino/tools/mo/utils/ir_reader/layer_to_class.py b/tools/mo/openvino/tools/mo/utils/ir_reader/layer_to_class.py index 0a5f393f9e3..618b9ccc7ce 100644 --- a/tools/mo/openvino/tools/mo/utils/ir_reader/layer_to_class.py +++ b/tools/mo/openvino/tools/mo/utils/ir_reader/layer_to_class.py @@ -32,6 +32,8 @@ from openvino.tools.mo.ops.split import Split, VariadicSplit from openvino.tools.mo.utils.class_registration import update_registration from openvino.tools.mo.utils.import_extensions import import_by_path from openvino.tools.mo.utils.ir_reader.extender import Extender +from openvino.tools.mo.utils.ir_reader.internal_ops.squeeze import SqueezeInternal +from openvino.tools.mo.utils.ir_reader.internal_ops.unsqueeze import UnsqueezeInternal # Operations not registered in collect_ops() function custom_ops = { @@ -50,9 +52,11 @@ custom_ops = { 'Power': Pow, 'Slice': OvSlice, 'Split': Split, + 'Squeeze': SqueezeInternal, 'Subtract': Sub, 'VariadicSplit': VariadicSplit, 'Clamp': AttributedClamp, + 'Unsqueeze': UnsqueezeInternal, } diff --git a/tools/mo/unit_tests/mo/utils/ir_reader/layer_to_class_test.py b/tools/mo/unit_tests/mo/utils/ir_reader/layer_to_class_test.py index 21a0e091301..f2a16bfc0a8 100644 --- a/tools/mo/unit_tests/mo/utils/ir_reader/layer_to_class_test.py +++ b/tools/mo/unit_tests/mo/utils/ir_reader/layer_to_class_test.py @@ -8,6 +8,8 @@ from generator import generator, generate from openvino.tools.mo.graph.graph import Node from openvino.tools.mo.utils.ir_engine.compare_graphs import compare_graphs +from openvino.tools.mo.utils.ir_reader.internal_ops.squeeze import SqueezeInternal +from openvino.tools.mo.utils.ir_reader.internal_ops.unsqueeze import UnsqueezeInternal from openvino.tools.mo.utils.ir_reader.layer_to_class import groupconv_to_conv, restore_tensor_names from unit_tests.utils.graph import build_graph @@ -128,5 +130,71 @@ class TestFunction(unittest.TestCase): assert node_1['fw_tensor_debug_info'] == [('abc', 'abc'), ('def', 'def')], 'Restored debug info is wrong!' assert node_2['fw_tensor_debug_info'] == [('ghi,jkl', 'ghi,jkl')], 'Restored debug info is wrong!' - assert node_3['fw_tensor_debug_info'] == [('mno', 'mno'), ('pqr,stu', 'pqr,stu')],\ + assert node_3['fw_tensor_debug_info'] == [('mno', 'mno'), ('pqr,stu', 'pqr,stu')], \ 'Restored debug info is wrong!' + + def test_squeeze(self): + nodes_attributes = { + 'input': {'kind': 'op', 'type': 'Parameter'}, + 'input_data': {'shape': [2, 1, 3], 'kind': 'data'}, + + 'axis': {'kind': 'op', 'type': 'Const', 'op': 'Const', 'value': np.array(1), 'shape': []}, + 'axis_data': {'shape': [], 'kind': 'data', 'value': np.array(1)}, + + 'squeeze': {'kind': 'op', 'type': 'Squeeze'}, + 'squeeze_data': {'shape': [2, 3], 'kind': 'data', 'value': None}, + + 'result': {'kind': 'op', 'type': 'Result'} + } + + edges = [('input', 'input_data'), + ('input_data', 'squeeze'), + ('axis', 'axis_data'), + ('axis_data', 'squeeze'), + ('squeeze', 'squeeze_data'), + ('squeeze_data', 'result'), + ] + + graph = build_graph(nodes_attributes, edges, nodes_with_edges_only=True) + + squeeze_node = Node(graph, 'squeeze') + SqueezeInternal.infer(squeeze_node) + + graph_ref = build_graph(nodes_attributes, edges, nodes_with_edges_only=True) + + # Check that graph wasn't changed after shape infer + (flag, resp) = compare_graphs(graph, graph_ref, 'result', check_op_attrs=True) + self.assertTrue(flag, resp) + + def test_unsqueeze(self): + nodes_attributes = { + 'input': {'kind': 'op', 'type': 'Parameter'}, + 'input_data': {'shape': [2, 3], 'kind': 'data'}, + + 'axis': {'kind': 'op', 'type': 'Const', 'op': 'Const', 'value': np.array(1), 'shape': []}, + 'axis_data': {'shape': [], 'kind': 'data', 'value': np.array(1)}, + + 'unsqueeze': {'kind': 'op', 'type': 'Unsqueeze'}, + 'unsqueeze_data': {'shape': [2, 1, 3], 'kind': 'data', 'value': None}, + + 'result': {'kind': 'op', 'type': 'Result'} + } + + edges = [('input', 'input_data'), + ('input_data', 'unsqueeze'), + ('axis', 'axis_data'), + ('axis_data', 'unsqueeze'), + ('unsqueeze', 'unsqueeze_data'), + ('unsqueeze_data', 'result'), + ] + + graph = build_graph(nodes_attributes, edges, nodes_with_edges_only=True) + + unsqueeze_node = Node(graph, 'unsqueeze') + UnsqueezeInternal.infer(unsqueeze_node) + + graph_ref = build_graph(nodes_attributes, edges, nodes_with_edges_only=True) + + # Check that graph wasn't changed after shape infer + (flag, resp) = compare_graphs(graph, graph_ref, 'result', check_op_attrs=True) + self.assertTrue(flag, resp)