[ MO ] Random Uniform Replacer (#1814)

diff --git a/model-optimizer/automation/package_BOM.txt b/model-optimizer/automation/package_BOM.txt
index 682cdee..7c98a0e 100644 (file)
--- a/model-optimizer/automation/package_BOM.txt
+++ b/model-optimizer/automation/package_BOM.txt
@@ -540,6 +540,7 @@ extensions/middle/pass_separator.py
 extensions/middle/permute_tensor_iterator.py
 extensions/middle/preprocessing.py
 extensions/middle/quantize_fuses.py
+extensions/middle/RandomUniformReplacer.py
 extensions/middle/ReluQuantizeFuse.py
 extensions/middle/RemoveDuplicationMemory.py
 extensions/middle/RemoveIdentity.py

diff --git a/model-optimizer/extensions/middle/RandomUniformReplacer.py b/model-optimizer/extensions/middle/RandomUniformReplacer.py
new file mode 100644 (file)
index 0000000..389aa04
--- /dev/null
+++ b/model-optimizer/extensions/middle/RandomUniformReplacer.py
@@ -0,0 +1,66 @@
+"""
+ Copyright (C) 2018-2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+import numpy as np
+
+from typing import Dict
+
+from mo.front.tf.graph_utils import create_op_with_const_inputs
+from mo.graph.graph import Graph, Node
+from mo.middle.replacement import MiddleReplacementPattern
+from mo.ops.broadcast import Broadcast
+
+
+class RandomUniformReplacer(MiddleReplacementPattern):
+    """
+    Replaces RandomUniform operation with Broadcast of ones in sub-graph:
+
+    ShapeOf ---> RandomUniform ---> Mul
+
+    """
+
+    enabled = True
+
+    @staticmethod
+    def pattern():
+        return dict(
+            nodes=[
+                ('shape', dict(op='ShapeOf')),
+                ('shape_data', dict()),
+                ('random_uniform', dict(op='RandomUniform')),
+                ('random_uniform_data', dict()),
+                ('mul', dict(op='Mul')),
+                ('mul_const', dict(op='Const')),
+                ('mul_const_data', dict())
+            ],
+            edges=[
+                ('shape', 'shape_data'),
+                ('shape_data', 'random_uniform'),
+                ('random_uniform', 'random_uniform_data'),
+                ('random_uniform_data', 'mul'),
+                ('mul_const', 'mul_const_data'),
+                ('mul_const_data', 'mul')
+            ]
+        )
+
+    @staticmethod
+    def replace_pattern(graph: Graph, match: Dict[str, Node]):
+        node = match['random_uniform']
+        node_name = node.soft_get('name', node.id)
+        data_type = match['mul_const'].out_port(0).get_data_type()
+        broadcast_node = create_op_with_const_inputs(graph, Broadcast, port_value_dict={0: np.array([1], dtype=data_type)},
+                                                     op_attrs={'name': node_name + '/Broadcast', 'mode': 'numpy'})
+        node.in_port(0).get_connection().set_destination(broadcast_node.in_port(1))
+        node.out_port(0).get_connection().set_source(broadcast_node.out_port(0))

diff --git a/model-optimizer/extensions/middle/RandomUniformReplacer_test.py b/model-optimizer/extensions/middle/RandomUniformReplacer_test.py
new file mode 100644 (file)
index 0000000..a73d231
--- /dev/null
+++ b/model-optimizer/extensions/middle/RandomUniformReplacer_test.py
@@ -0,0 +1,98 @@
+"""
+ Copyright (C) 2018-2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+import unittest
+
+import numpy as np
+
+from extensions.middle.RandomUniformReplacer import RandomUniformReplacer
+from mo.utils.ir_engine.compare_graphs import compare_graphs
+from mo.utils.unittest.graph import build_graph
+
+nodes_attributes = {
+    'shape': {'kind': 'op', 'op': 'ShapeOf'},
+    'shape_data': {'kind': 'data'},
+    'random_uniform': {'kind': 'op', 'op': 'RandomUniform'},
+    'random_uniform_data': {'kind': 'data'},
+    'mul': {'kind': 'op', 'op': 'Mul'},
+    'mul_const': {'kind': 'op', 'op': 'Const'},
+    'mul_const_data': {'kind': 'data', 'value': np.array([1], dtype=np.int32)},
+
+    'broadcast': {'kind': 'op', 'op': 'Broadcast'},
+    'broadcast_const': {'kind': 'op', 'op': 'Const'},
+    'broadcast_const_data': {'kind': 'data', 'value': np.array([1], dtype=np.int32)},
+}
+
+
+class RandomUniformReplacerTest(unittest.TestCase):
+    def test_1(self):
+        graph = build_graph(nodes_attributes,
+                            edges=[
+                                ('shape', 'shape_data'),
+                                ('shape_data', 'random_uniform'),
+                                ('random_uniform', 'random_uniform_data'),
+                                ('random_uniform_data', 'mul', {'in': 0}),
+                                ('mul_const', 'mul_const_data'),
+                                ('mul_const_data', 'mul', {'in': 1})
+                            ],
+                            nodes_with_edges_only=True)
+
+        ref_graph = build_graph(nodes_attributes,
+                                edges=[
+                                    ('shape', 'shape_data'),
+                                    ('shape_data', 'broadcast', {'in': 1}),
+                                    ('broadcast_const', 'broadcast_const_data'),
+                                    ('broadcast_const_data', 'broadcast', {'in': 0}),
+                                    ('broadcast', 'random_uniform_data'),
+                                    ('random_uniform_data', 'mul', {'in': 0}),
+                                    ('mul_const', 'mul_const_data'),
+                                    ('mul_const_data', 'mul', {'in': 1})
+                                ],
+                                nodes_with_edges_only=True)
+
+        RandomUniformReplacer().find_and_replace_pattern(graph)
+
+        flag, resp = compare_graphs(graph, ref_graph, 'mul', check_op_attrs=True)
+        self.assertTrue(flag, resp)
+
+    def test_2(self):
+        graph = build_graph(nodes_attributes,
+                            edges=[
+                                ('shape', 'shape_data'),
+                                ('shape_data', 'random_uniform'),
+                                ('random_uniform', 'random_uniform_data'),
+                                ('random_uniform_data', 'mul', {'in': 1}),
+                                ('mul_const', 'mul_const_data'),
+                                ('mul_const_data', 'mul', {'in': 0})
+                            ],
+                            nodes_with_edges_only=True)
+
+        ref_graph = build_graph(nodes_attributes,
+                                edges=[
+                                    ('shape', 'shape_data'),
+                                    ('shape_data', 'broadcast', {'in': 1}),
+                                    ('broadcast_const', 'broadcast_const_data'),
+                                    ('broadcast_const_data', 'broadcast', {'in': 0}),
+                                    ('broadcast', 'random_uniform_data'),
+                                    ('random_uniform_data', 'mul', {'in': 1}),
+                                    ('mul_const', 'mul_const_data'),
+                                    ('mul_const_data', 'mul', {'in': 0})
+                                ],
+                                nodes_with_edges_only=True)
+
+        RandomUniformReplacer().find_and_replace_pattern(graph)
+
+        flag, resp = compare_graphs(graph, ref_graph, 'mul', check_op_attrs=True)
+        self.assertTrue(flag, resp)