Extend GELU merger patterns (#4130)

This commit is contained in:
Maxim Vafin authored and GitHub committed 2021-02-02 17:59:12 +03:00
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')