Extend LayerNorm pattern to support onnx bert (#4137)
* Extend LayerNorm pattern to support onnx bert * Fix test * Fix test
This commit is contained in:
Maxim Vafin authored and 
GitHub committed
1 parent
65e2b4af3f
commit
08a527602f
2 files changed
+105
-39
No files matched your search
@@ -16,18 +16,18 @@
|
||||
|
||||
import logging as log
|
||||
|
||||
from mo.front.common.replacement import FrontReplacementSubgraph
|
||||
from mo.front.common.replacement import FrontReplacementPattern
|
||||
from mo.front.tf.graph_utils import create_op_with_const_inputs
|
||||
from mo.graph.graph import Graph, rename_nodes
|
||||
from extensions.ops.mvn import MVN
|
||||
from mo.middle.pattern_match import apply_pattern
|
||||
|
||||
|
||||
class LayerNorm(FrontReplacementSubgraph):
|
||||
class LayerNorm(FrontReplacementPattern):
|
||||
# Compose part of the LayerNorm pattern to the MVN
|
||||
enabled = True
|
||||
|
||||
def pattern(self):
|
||||
log.info('Enabled LayerNorm pattern recognition')
|
||||
def pattern1(self):
|
||||
return dict(
|
||||
nodes=[
|
||||
('pool0', dict(op='ReduceMean')),
|
||||
@@ -56,21 +56,59 @@ class LayerNorm(FrontReplacementSubgraph):
|
||||
('add_param', 'add'),
|
||||
])
|
||||
|
||||
def replace_sub_graph(self, graph: Graph, match: dict):
|
||||
def pattern2(self):
|
||||
# pattern from bert onnx model
|
||||
return dict(
|
||||
nodes=[
|
||||
('pool0', dict(op='ReduceMean')),
|
||||
('pool1', dict(op='ReduceMean')),
|
||||
('cast', dict(op='Cast')),
|
||||
('pow', dict(op='Pow')),
|
||||
('div', dict(op='Div')),
|
||||
('sqrt', dict(op='Pow')),
|
||||
('add', dict(op='Add')),
|
||||
('sub', dict(op='Sub')),
|
||||
('pool0_param', dict(op='Const')),
|
||||
('pool1_param', dict(op='Const')),
|
||||
('add_param', dict(op='Const')),
|
||||
('pow_param', dict(op='Const')),
|
||||
],
|
||||
edges=[
|
||||
('pool0', 'sub'),
|
||||
('sub', 'cast'),
|
||||
('cast', 'pow'),
|
||||
('pow', 'pool1'),
|
||||
('pool1', 'add'),
|
||||
('add', 'sqrt'),
|
||||
('sqrt', 'div'),
|
||||
('sub', 'div'),
|
||||
('pool0_param', 'pool0'),
|
||||
('pool1_param', 'pool1'),
|
||||
('pow_param', 'sqrt'),
|
||||
('add_param', 'add'),
|
||||
])
|
||||
|
||||
def find_and_replace_pattern(self, graph: Graph):
|
||||
log.info('Enabled LayerNorm pattern recognition')
|
||||
apply_pattern(graph, **self.pattern1(), action=self.replace_layer_norm)
|
||||
apply_pattern(graph, **self.pattern2(), action=self.replace_layer_norm)
|
||||
|
||||
def replace_layer_norm(self, graph: Graph, match: dict):
|
||||
inp = match['pool0']
|
||||
inp_port = inp.in_port(0).get_source()
|
||||
node_before = inp.in_port(0).get_source().node
|
||||
node_before_name = node_before.soft_get('name', node_before.id)
|
||||
|
||||
# take/check the values of the add, pow and axes for ReduceMean
|
||||
pow_param = match['pow_param']
|
||||
add_param = match['add_param']
|
||||
if add_param.value.size == 1 and pow_param.value.size == 1 and add_param.value.item() <= 1e-05 \
|
||||
and pow_param.value.item() == 0.5 and match['pool0_param'].value == match['pool1_param'].value:
|
||||
log.debug('Found LayerNorm pattern after {} with name {}'.format(inp_port.node.op, inp_port.node.name))
|
||||
log.debug('Found LayerNorm pattern after {} with name {}'.format(node_before.op, node_before_name))
|
||||
mvn = create_op_with_const_inputs(graph, MVN, {1: match['pool1_param'].value},
|
||||
{'eps': add_param.value.item(), 'normalize_variance': 1,
|
||||
'eps_mode': 'inside_sqrt'})
|
||||
div_name = match['div'].soft_get('name', match['div'].id)
|
||||
rename_nodes([(match['div'], div_name + '/to_be_removed'), (mvn, div_name)])
|
||||
|
||||
inp_port.connect(mvn.in_port(0))
|
||||
inp.in_port(0).get_connection().set_destination(mvn.in_port(0))
|
||||
match['div'].out_port(0).get_connection().set_source(mvn.out_port(0))
|
||||
@@ -22,33 +22,34 @@ from extensions.front.LayerNorm import LayerNorm
|
||||
from mo.utils.ir_engine.compare_graphs import compare_graphs
|
||||
from mo.utils.unittest.graph import build_graph
|
||||
|
||||
nodes_attributes_mvn = {
|
||||
'inp': {'kind': 'op', 'op': 'AnyOp'},
|
||||
'pool0': {'kind': 'op', 'op': 'ReduceMean'},
|
||||
'pool1': {'kind': 'op', 'op': 'ReduceMean'},
|
||||
'pow': {'kind': 'op', 'op': 'Pow'},
|
||||
'div': {'kind': 'op', 'op': 'Div'},
|
||||
'sqrt': {'kind': 'op', 'op': 'Pow'},
|
||||
'add': {'kind': 'op', 'op': 'Add'},
|
||||
'sub': {'kind': 'op', 'op': 'Sub'},
|
||||
'add_param': {'kind': 'op', 'op': 'Const'},
|
||||
'pow_param': {'kind': 'op', 'op': 'Const'},
|
||||
'pool0_param': {'kind': 'op', 'op': 'Const'},
|
||||
'pool1_param': {'kind': 'op', 'op': 'Const'},
|
||||
'out': {'kind': 'op', 'op': 'AnyOp'},
|
||||
}
|
||||
|
||||
nodes_attributes_ref = {
|
||||
'inp': {'kind': 'op', 'op': 'AnyOp'},
|
||||
'mvn': {'kind': 'op', 'op': 'MVN', 'eps': 1e-6, 'normalize_variance': 1, 'eps_mode': 'inside_sqrt'},
|
||||
'mvn_param': {'kind': 'op', 'op': 'Const'},
|
||||
'out': {'kind': 'op', 'op': 'AnyOp'},
|
||||
}
|
||||
|
||||
|
||||
class TestMVNPatternReplacement(unittest.TestCase):
|
||||
nodes_attributes_mvn = {
|
||||
'inp': {'kind': 'op', 'op': 'AnyOp'},
|
||||
'pool0': {'kind': 'op', 'op': 'ReduceMean'},
|
||||
'pool1': {'kind': 'op', 'op': 'ReduceMean'},
|
||||
'cast': {'kind': 'op', 'op': 'Cast'},
|
||||
'pow': {'kind': 'op', 'op': 'Pow'},
|
||||
'div': {'kind': 'op', 'op': 'Div'},
|
||||
'sqrt': {'kind': 'op', 'op': 'Pow'},
|
||||
'add': {'kind': 'op', 'op': 'Add'},
|
||||
'sub': {'kind': 'op', 'op': 'Sub'},
|
||||
'add_param': {'kind': 'op', 'op': 'Const', 'shape': np.array([1]), 'value': np.array(1e-06, dtype=np.float32)},
|
||||
'pow_param': {'kind': 'op', 'op': 'Const', 'shape': np.array([1]), 'value': np.array(0.5, dtype=np.float32)},
|
||||
'pool0_param': {'kind': 'op', 'op': 'Const', 'shape': np.array([1]), 'value': np.array(-1, dtype=np.int32)},
|
||||
'pool1_param': {'kind': 'op', 'op': 'Const', 'shape': np.array([1]), 'value': np.array(-1, dtype=np.int32)},
|
||||
'out': {'kind': 'op', 'op': 'AnyOp'},
|
||||
}
|
||||
|
||||
nodes_attributes_ref = {
|
||||
'inp': {'kind': 'op', 'op': 'AnyOp'},
|
||||
'mvn': {'kind': 'op', 'op': 'MVN', 'eps': 1e-6, 'normalize_variance': 1, 'eps_mode': 'inside_sqrt'},
|
||||
'mvn_param': {'kind': 'op', 'op': 'Const', 'shape': np.array([]), 'value': np.array(-1, dtype=np.int32)},
|
||||
'out': {'kind': 'op', 'op': 'AnyOp'},
|
||||
}
|
||||
|
||||
def test_MVNPatternReplacement_test_1(self):
|
||||
graph = build_graph(nodes_attributes_mvn,
|
||||
graph = build_graph(self.nodes_attributes_mvn,
|
||||
[('inp', 'pool0', {'out': 0}),
|
||||
('inp', 'sub', {'out': 0}),
|
||||
('pool0', 'sub'),
|
||||
@@ -64,17 +65,44 @@ class TestMVNPatternReplacement(unittest.TestCase):
|
||||
('pool0_param', 'pool0'),
|
||||
('pool1_param', 'pool1'),
|
||||
],
|
||||
{'pow_param': {'shape': np.array([1]), 'value': np.array(0.5)},
|
||||
'add_param': {'shape': np.array([1]), 'value': np.array(1e-06)},
|
||||
'pool0_param': {'shape': np.array([1]), 'value': np.array(-1)},
|
||||
'pool1_param': {'shape': np.array([1]), 'value': np.array(-1)},
|
||||
},
|
||||
nodes_with_edges_only=True)
|
||||
graph_ref = build_graph(nodes_attributes_ref,
|
||||
graph_ref = build_graph(self.nodes_attributes_ref,
|
||||
[('inp', 'mvn'),
|
||||
('mvn_param', 'mvn'),
|
||||
('mvn', 'out')],
|
||||
{}, nodes_with_edges_only=True)
|
||||
nodes_with_edges_only=True)
|
||||
graph.stage = 'front'
|
||||
|
||||
replacer = LayerNorm()
|
||||
replacer.find_and_replace_pattern(graph)
|
||||
|
||||
(flag, resp) = compare_graphs(graph, graph_ref, 'out', check_op_attrs=True)
|
||||
self.assertTrue(flag, resp)
|
||||
|
||||
def test_MVNPatternReplacement_test_2(self):
|
||||
graph = build_graph(self.nodes_attributes_mvn,
|
||||
[('inp', 'pool0', {'out': 0}),
|
||||
('inp', 'sub', {'out': 0}),
|
||||
('pool0', 'sub'),
|
||||
('sub', 'cast'),
|
||||
('cast', 'pow'),
|
||||
('pow', 'pool1'),
|
||||
('pool1', 'add'),
|
||||
('add', 'sqrt'),
|
||||
('sqrt', 'div'),
|
||||
('sub', 'div'),
|
||||
('div', 'out'),
|
||||
('pow_param', 'sqrt'),
|
||||
('add_param', 'add'),
|
||||
('pool0_param', 'pool0'),
|
||||
('pool1_param', 'pool1'),
|
||||
],
|
||||
nodes_with_edges_only=True)
|
||||
graph_ref = build_graph(self.nodes_attributes_ref,
|
||||
[('inp', 'mvn'),
|
||||
('mvn_param', 'mvn'),
|
||||
('mvn', 'out')],
|
||||
nodes_with_edges_only=True)
|
||||
graph.stage = 'front'
|
||||
|
||||
replacer = LayerNorm()
|
||||
|
||||
Reference in new issue
Block a user