diff --git a/model-optimizer/automation/package_BOM.txt b/model-optimizer/automation/package_BOM.txt index be0dc1c37e0..e7c779543cd 100644 --- a/model-optimizer/automation/package_BOM.txt +++ b/model-optimizer/automation/package_BOM.txt @@ -57,6 +57,7 @@ extensions/back/ReverseInputChannels.py extensions/back/RNNSequenceTypeRename.py extensions/back/ScalarConstNormalize.py extensions/back/SelectBroadcast.py +extensions/back/ShapeOfConstFolding.py extensions/back/ShuffleChannelPatternOptimization.py extensions/back/ShufflenetReLUReorder.py extensions/back/SpecialNodesFinalization.py diff --git a/model-optimizer/extensions/back/ShapeOfConstFolding.py b/model-optimizer/extensions/back/ShapeOfConstFolding.py new file mode 100644 index 00000000000..b97b46fc35d --- /dev/null +++ b/model-optimizer/extensions/back/ShapeOfConstFolding.py @@ -0,0 +1,29 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from mo.back.replacement import BackReplacementPattern +from mo.graph.graph import Graph, rename_nodes +from mo.ops.const import Const + + +class ShapeOfConstFolding(BackReplacementPattern): + """ + The transformation folds ShapeOf(Const) -> Const + """ + enabled = True + + def run_after(self): + from extensions.back.MatMulNormalizer import SmartReshape_HC_Reshape_MatMul + return [SmartReshape_HC_Reshape_MatMul] + + def find_and_replace_pattern(self, graph: Graph): + for shapeof_node in graph.get_op_nodes(op='ShapeOf'): + in_node = shapeof_node.in_port(0).get_source().node + if in_node.op == 'Const': + shapeof_node.in_port(0).disconnect() + shape_name = shapeof_node.soft_get('name', shapeof_node.id) + shape_value = shapeof_node.out_port(0).data.get_value() + shape_const_node = Const(graph, {'name': shape_name + '/ExecutionConstValue', + 'value': shape_value}).create_node() + shapeof_node.out_port(0).get_connection().set_source(shape_const_node.out_port(0)) + rename_nodes([(shapeof_node, shape_name + '/TBD'), (shape_const_node, shape_name)]) diff --git a/model-optimizer/unit_tests/extensions/back/ShapeOfConstFolding_test.py b/model-optimizer/unit_tests/extensions/back/ShapeOfConstFolding_test.py new file mode 100644 index 00000000000..562c1416342 --- /dev/null +++ b/model-optimizer/unit_tests/extensions/back/ShapeOfConstFolding_test.py @@ -0,0 +1,170 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +import numpy as np + +from extensions.back.ShapeOfConstFolding import ShapeOfConstFolding +from mo.front.common.partial_infer.eltwise import eltwise_infer +from mo.front.common.partial_infer.utils import int64_array +from mo.utils.ir_engine.compare_graphs import compare_graphs +from unit_tests.utils.graph import build_graph + +const_value = np.random.rand(1, 3, 30, 30) +nodes_attributes = {'input': {'shape': int64_array([1, 3, 30, 30]), 'type': 'Parameter', 'kind': 'op', + 'op': 'Parameter'}, + 'input_data': {'value': None, 'shape': int64_array([1, 3, 30, 30]), 'kind': 'data'}, + 'const': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': const_value}, + 'const_data': {'kind': 'data', 'value': const_value}, + 'shapeof_input': {'kind': 'op', 'op': 'ShapeOf', 'value': int64_array([1, 3, 30, 30])}, + 'shapeof_input_data': {'value': None, 'shape': None, 'kind': 'data', + 'value': int64_array([1, 3, 30, 30])}, + + 'shapeof_const': {'kind': 'op', 'op': 'ShapeOf', 'value': int64_array([1, 3, 30, 30])}, + 'shapeof_const_data': {'value': None, 'shape': None, 'kind': 'data', + 'value': int64_array([1, 3, 30, 30])}, + + 'mul': {'kind': 'op', 'op': 'Mul', 'infer': lambda node: eltwise_infer(node, lambda a, b: a * b)}, + 'mul_data': {'kind': 'data', 'value': np.array([1, 9, 900, 900])}, + 'last': {'kind': 'op', 'op': 'Result'}, + + # new nodes + 'new_const_shapeof': {'type': 'Const', 'kind': 'op', 'op': 'Const', + 'value': int64_array([1, 3, 30, 30])} + } + +const_value2 = np.random.rand(30, 30) +nodes_attributes2 = {'input': {'shape': int64_array([1, 3, 30, 30]), 'type': 'Parameter', 'kind': 'op', + 'op': 'Parameter'}, + 'input_data': {'value': None, 'shape': int64_array([1, 3, 30, 30]), 'kind': 'data'}, + + 'const': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': const_value2}, + 'const_data': {'kind': 'data', 'value': const_value2}, + + 'shapeof_const': {'kind': 'op', 'op': 'ShapeOf', 'value': int64_array([2700, 30])}, + 'shapeof_const_data': {'value': int64_array([2700, 30]), 'shape': None, 'kind': 'data'}, + + 'gather': {'kind': 'op', 'op': 'Gather', 'batch_dims': 0}, + 'gather_data': {'kind': 'data'}, + + 'const_concat': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': [1]}, + 'const_concat_data': {'kind': 'data', 'value': [1]}, + 'concat': {'kind': 'op', 'op': 'Concat'}, + 'concat_data': {'kind': 'data'}, + + 'reshape': {'kind': 'op', 'op': 'Reshape'}, + 'reshape_data': {'kind': 'data'}, + + 'matmul': {'kind': 'op', 'op': 'MatMul'}, + 'matmul_data': {'kind': 'data'}, + 'last': {'kind': 'op', 'op': 'Result'}, + + # new nodes + 'new_const_shapeof': {'type': 'Const', 'kind': 'op', 'op': 'Const', + 'value': int64_array([2700, 30])}, + } + + +class ShapeOfConstFoldingTests(unittest.TestCase): + def test_const_with_one_output(self): + graph = build_graph(nodes_attributes, + [('input', 'input_data'), + ('input_data', 'shapeof_input'), + ('shapeof_input', 'shapeof_input_data'), + ('shapeof_input_data', 'mul'), + ('const', 'const_data'), + ('const_data', 'shapeof_const'), + ('shapeof_const', 'shapeof_const_data'), + ('shapeof_const_data', 'mul'), + ('mul', 'mul_data'), + ('mul_data', 'last')], + { + 'input': {'shape': int64_array([1, 3, 30, 30])}, + 'input_data': {'shape': int64_array([1, 3, 30, 30])}, + 'shapeof_input': {'value': int64_array([1, 3, 30, 30])}, + 'shapeof_input_data': {'value': int64_array([1, 3, 30, 30])}, + 'const': {'value': const_value}, + 'const_data': {'value': const_value}, + 'shapeof_const': {'value': int64_array([1, 3, 30, 30])}, + 'shapeof_const_data': {'value': int64_array([1, 3, 30, 30])}, + 'mul_data': {'value': int64_array([1, 9, 900, 900])}, + }, + nodes_with_edges_only=True) + + graph_ref = build_graph(nodes_attributes, + [('input', 'input_data'), + ('input_data', 'shapeof_input'), + ('shapeof_input', 'shapeof_input_data'), + ('shapeof_input_data', 'mul'), + ('new_const_shapeof', 'shapeof_const_data'), + ('shapeof_const_data', 'mul'), + ('mul', 'mul_data'), + ('mul_data', 'last')], + { + 'input': {'shape': int64_array([1, 3, 30, 30])}, + 'input_data': {'shape': int64_array([1, 3, 30, 30])}, + 'shapeof_input': {'value': int64_array([1, 3, 30, 30])}, + 'shapeof_input_data': {'value': int64_array([1, 3, 30, 30])}, + 'new_const_shapeof': {'value': int64_array([1, 3, 30, 30])}, + 'shapeof_const_data': {'value': int64_array([1, 3, 30, 30])}, + 'mul_data': {'value': int64_array([1, 9, 900, 900])}, + }, + nodes_with_edges_only=True) + ShapeOfConstFolding().find_and_replace_pattern(graph) + (flag, resp) = compare_graphs(graph, graph_ref, 'last') + self.assertTrue(flag, resp) + + def test_const_with_two_outputs(self): + graph = build_graph(nodes_attributes2, + [('input', 'input_data'), + ('input_data', 'reshape'), + ('const', 'const_data'), + ('const_data', 'shapeof_const'), + ('shapeof_const', 'shapeof_const_data'), + ('shapeof_const_data', 'gather'), + ('gather', 'gather_data'), + ('const_concat', 'const_concat_data'), + ('const_concat_data', 'concat'), + ('gather_data', 'concat'), + ('concat', 'reshape'), + ('reshape', 'reshape_data'), + ('reshape_data', 'matmul'), + ('const_data', 'matmul'), + ('matmul', 'matmul_data'), + ('matmul_data', 'last') + ], + { + 'input': {'shape': int64_array([1, 3, 30, 30])}, + 'input_data': {'shape': int64_array([1, 3, 30, 30])}, + 'shapeof_const': {'value': int64_array([2700, 30])}, + 'shapeof_const_data': {'value': int64_array([2700, 30])}, + }, + nodes_with_edges_only=True) + + graph_ref = build_graph(nodes_attributes2, + [('input', 'input_data'), + ('input_data', 'reshape'), + ('new_const_shapeof', 'shapeof_const_data'), + ('shapeof_const_data', 'gather'), + ('gather', 'gather_data'), + ('const_concat', 'const_concat_data'), + ('const_concat_data', 'concat'), + ('gather_data', 'concat'), + ('concat', 'reshape'), + ('reshape', 'reshape_data'), + ('reshape_data', 'matmul'), + ('const', 'const_data'), + ('const_data', 'matmul'), + ('matmul', 'matmul_data'), + ('matmul_data', 'last')], + { + 'input': {'shape': int64_array([1, 3, 30, 30])}, + 'input_data': {'shape': int64_array([1, 3, 30, 30])}, + 'new_const_shapeof': {'value': int64_array([2700, 30])}, + 'shapeof_const_data': {'value': int64_array([2700, 30])}, + }, + nodes_with_edges_only=True) + ShapeOfConstFolding().find_and_replace_pattern(graph) + (flag, resp) = compare_graphs(graph, graph_ref, 'last') + self.assertTrue(flag, resp)