Extend GELU merger patterns (#4130)
This commit is contained in:
1 parent
5c08312ade
commit
75763b8e67
2 files changed
+191
-19
No files matched your search
@@ -1,5 +1,5 @@
|
||||
"""
|
||||
Copyright (C) 2017-2020 Intel Corporation
|
||||
Copyright (C) 2017-2021 Intel Corporation
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
@@ -19,26 +19,26 @@ import logging as log
|
||||
from math import sqrt, fabs
|
||||
|
||||
from extensions.ops.gelu import GeLUOP
|
||||
from mo.front.common.replacement import FrontReplacementSubgraph
|
||||
from mo.graph.graph import Graph
|
||||
from mo.front.common.replacement import FrontReplacementPattern
|
||||
from mo.graph.graph import Graph, rename_nodes
|
||||
from mo.middle.pattern_match import apply_pattern
|
||||
|
||||
|
||||
class GeLUMergerErf(FrontReplacementSubgraph):
|
||||
class GeLUMergerErf(FrontReplacementPattern):
|
||||
enabled = True
|
||||
|
||||
def pattern(self):
|
||||
log.info('Enabled GeLU Merger replacement for approximation with Erf')
|
||||
def pattern1(self):
|
||||
# (0.5 * x) * (1 + erf(x / sqrt(2))
|
||||
return dict(
|
||||
nodes=[
|
||||
('mul', dict(op='Mul')),
|
||||
('mul', dict(op='Mul')),
|
||||
('mul0', dict(op='Mul')),
|
||||
('div', dict(op='Div')),
|
||||
('erf', dict(op='Erf')),
|
||||
('add', dict(op='Add')),
|
||||
('div', dict(op='Div')),
|
||||
('erf', dict(op='Erf')),
|
||||
('add', dict(op='Add')),
|
||||
('mul_param', dict(op='Const')),
|
||||
('div_param', dict(op='Const')),
|
||||
('add_param', dict(op='Const')),
|
||||
|
||||
],
|
||||
edges=[
|
||||
('mul', 'mul0'),
|
||||
@@ -50,13 +50,67 @@ class GeLUMergerErf(FrontReplacementSubgraph):
|
||||
('add_param', 'add'),
|
||||
])
|
||||
|
||||
def replace_sub_graph(self, graph: Graph, match: dict):
|
||||
def pattern2(self):
|
||||
# 0.5 * (x * (1 + erf(x / sqrt(2)))
|
||||
return dict(
|
||||
nodes=[
|
||||
('mul', dict(op='Mul')),
|
||||
('mul0', dict(op='Mul')),
|
||||
('div', dict(op='Div')),
|
||||
('erf', dict(op='Erf')),
|
||||
('add', dict(op='Add')),
|
||||
('mul_param', dict(op='Const')),
|
||||
('div_param', dict(op='Const')),
|
||||
('add_param', dict(op='Const')),
|
||||
],
|
||||
edges=[
|
||||
('div', 'erf'),
|
||||
('erf', 'add'),
|
||||
('add', 'mul'),
|
||||
('mul', 'mul0'),
|
||||
('mul_param', 'mul0'),
|
||||
('div_param', 'div'),
|
||||
('add_param', 'add'),
|
||||
])
|
||||
|
||||
def pattern3(self):
|
||||
# x * (0.5 * (1 + erf(x / sqrt(2)))
|
||||
return dict(
|
||||
nodes=[
|
||||
('mul', dict(op='Mul')),
|
||||
('mul0', dict(op='Mul')),
|
||||
('div', dict(op='Div')),
|
||||
('erf', dict(op='Erf')),
|
||||
('add', dict(op='Add')),
|
||||
('mul_param', dict(op='Const')),
|
||||
('div_param', dict(op='Const')),
|
||||
('add_param', dict(op='Const')),
|
||||
],
|
||||
edges=[
|
||||
('div', 'erf'),
|
||||
('erf', 'add'),
|
||||
('add', 'mul'),
|
||||
('mul', 'mul0'),
|
||||
('mul_param', 'mul'),
|
||||
('div_param', 'div'),
|
||||
('add_param', 'add'),
|
||||
])
|
||||
|
||||
def find_and_replace_pattern(self, graph: Graph):
|
||||
log.info('Enabled GeLU Merger replacement for approximation with Erf')
|
||||
apply_pattern(graph, **self.pattern1(), action=self.replace_gelu)
|
||||
apply_pattern(graph, **self.pattern2(), action=self.replace_gelu)
|
||||
apply_pattern(graph, **self.pattern3(), action=self.replace_gelu)
|
||||
|
||||
def replace_gelu(self, graph: Graph, match: dict):
|
||||
# Gaussian Error Linear Unit
|
||||
# f(x) = 0.5 * x * (1 + erf(x / sqrt(2))
|
||||
out_node = match['mul0']
|
||||
node_name = out_node.soft_get('name', out_node.id)
|
||||
div = match['div']
|
||||
inp_port = div.in_port(0).get_source()
|
||||
inp = inp_port.node
|
||||
log.debug('Found potential Erf-based GeLU pattern after {} with name {}'.format(inp.op, inp.name))
|
||||
inp_node = div.in_port(0).get_source().node
|
||||
inp_name = inp_node.soft_get('name', out_node.id)
|
||||
log.debug('Found potential Erf-based GeLU pattern after {} with name {}'.format(inp_node.op, inp_name))
|
||||
|
||||
# take the values of the mul, add and div
|
||||
div_param = match['div_param']
|
||||
@@ -71,7 +125,8 @@ class GeLUMergerErf(FrontReplacementSubgraph):
|
||||
sqrt2 = sqrt(2.0)
|
||||
# check that the values match the approximation
|
||||
if fabs(div_param - sqrt2) < 1e-06 and mul_param == 0.5 and add_param == 1.0:
|
||||
log.debug('Confirmed Erf-based GELU pattern after {} with name {}'.format(inp.op, inp.name))
|
||||
gelu = GeLUOP(graph, dict(name=inp.name + '/GELU_')).create_node()
|
||||
inp_port.connect(gelu.in_port(0))
|
||||
match['mul0'].out_port(0).get_connection().set_source(gelu.out_port(0))
|
||||
log.debug('Confirmed Erf-based GELU pattern after {} with name {}'.format(inp_node.op, inp_name))
|
||||
gelu = GeLUOP(graph, dict(name=inp_name + '/GELU_')).create_node()
|
||||
div.in_port(0).get_connection().set_destination(gelu.in_port(0))
|
||||
out_node.out_port(0).get_connection().set_source(gelu.out_port(0))
|
||||
rename_nodes([(out_node, node_name + '/TBD'), (gelu, node_name)])
|
||||
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
Copyright (C) 2018-2021 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
|
||||
from math import sqrt
|
||||
|
||||
from extensions.front.GeLUMerger_Erf import GeLUMergerErf
|
||||
from mo.front.common.partial_infer.utils import float_array, int64_array
|
||||
from mo.utils.ir_engine.compare_graphs import compare_graphs
|
||||
from mo.utils.unittest.graph import build_graph, const, regular_op, result, build_graph
|
||||
|
||||
ref_nodes = {**regular_op('input', {'type': 'Parameter'}),
|
||||
**regular_op('gelu', {'type': 'Gelu', 'name': 'final_mul'}),
|
||||
**result('result')
|
||||
}
|
||||
ref_edges = [('input', 'gelu'), ('gelu', 'result')]
|
||||
|
||||
|
||||
class GeLUMergerErfTest(unittest.TestCase):
|
||||
nodes = {
|
||||
**regular_op('input', {'op': 'Parameter', 'type': 'Parameter'}),
|
||||
**regular_op('mul', {'op': 'Mul'}),
|
||||
**regular_op('mul0', {'op': 'Mul', 'name': 'final_mul'}),
|
||||
**regular_op('div', {'op': 'Div'}),
|
||||
**regular_op('erf', {'op': 'Erf'}),
|
||||
**regular_op('add', {'op': 'Add'}),
|
||||
**const('mul_param', float_array([0.5])),
|
||||
**const('div_param', float_array([sqrt(2.)])),
|
||||
**const('add_param', int64_array([1])),
|
||||
**result('result'),
|
||||
}
|
||||
|
||||
def test_gelu_p1(self):
|
||||
edges = [('input', 'mul'),
|
||||
('mul', 'mul0'),
|
||||
('input', 'div'),
|
||||
('div', 'erf'),
|
||||
('erf', 'add'),
|
||||
('add', 'mul0'),
|
||||
('mul_param', 'mul'),
|
||||
('div_param', 'div'),
|
||||
('add_param', 'add'),
|
||||
('mul0', 'result')]
|
||||
|
||||
graph = build_graph(self.nodes, edges)
|
||||
|
||||
graph_ref = build_graph(ref_nodes, ref_edges)
|
||||
graph.stage = 'front'
|
||||
|
||||
GeLUMergerErf().find_and_replace_pattern(graph)
|
||||
graph.clean_up()
|
||||
|
||||
(flag, resp) = compare_graphs(graph, graph_ref, 'result')
|
||||
self.assertTrue(flag, resp)
|
||||
self.assertTrue(len(graph.get_op_nodes(name='final_mul')) == 1 and
|
||||
graph.get_op_nodes(name='final_mul')[0].op == 'Gelu')
|
||||
|
||||
def test_gelu_p2(self):
|
||||
edges = [('input', 'mul'),
|
||||
('div', 'erf'),
|
||||
('erf', 'add'),
|
||||
('add', 'mul'),
|
||||
('mul', 'mul0'),
|
||||
('mul_param', 'mul0'),
|
||||
('div_param', 'div'),
|
||||
('add_param', 'add'),
|
||||
('mul0', 'result')]
|
||||
|
||||
graph = build_graph(self.nodes, edges)
|
||||
|
||||
graph_ref = build_graph(ref_nodes, ref_edges)
|
||||
graph.stage = 'front'
|
||||
|
||||
GeLUMergerErf().find_and_replace_pattern(graph)
|
||||
graph.clean_up()
|
||||
|
||||
(flag, resp) = compare_graphs(graph, graph_ref, 'result')
|
||||
self.assertTrue(flag, resp)
|
||||
self.assertTrue(len(graph.get_op_nodes(name='final_mul')) == 1 and
|
||||
graph.get_op_nodes(name='final_mul')[0].op == 'Gelu')
|
||||
|
||||
def test_gelu_p3(self):
|
||||
edges = [('input', 'mul'),
|
||||
('div', 'erf'),
|
||||
('erf', 'add'),
|
||||
('add', 'mul'),
|
||||
('mul', 'mul0'),
|
||||
('mul_param', 'mul'),
|
||||
('div_param', 'div'),
|
||||
('add_param', 'add'),
|
||||
('mul0', 'result')]
|
||||
|
||||
graph = build_graph(self.nodes, edges)
|
||||
|
||||
graph_ref = build_graph(ref_nodes, ref_edges)
|
||||
graph.stage = 'front'
|
||||
|
||||
GeLUMergerErf().find_and_replace_pattern(graph)
|
||||
graph.clean_up()
|
||||
|
||||
(flag, resp) = compare_graphs(graph, graph_ref, 'result')
|
||||
self.assertTrue(flag, resp)
|
||||
self.assertTrue(len(graph.get_op_nodes(name='final_mul')) == 1 and
|
||||
graph.get_op_nodes(name='final_mul')[0].op == 'Gelu')
|
||||
Reference in new issue
Block a user