Add LogSoftmax-5 to MO and ngraph (#2409)
Co-authored-by: Evgeny Lazarev <evgeny.lazarev@intel.com>
This commit is contained in:
Maxim VafinandEvgeny Lazarev authored and 
GitHub committed
1 parent
83670dd5cb
commit
a405546054
23 files changed
+847
-240
No files matched your search
@@ -1,92 +0,0 @@
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
from extensions.ops.ReduceOps import ReduceMax, ReduceSum
|
||||
from extensions.ops.activation_ops import Exp, Log
|
||||
from extensions.ops.elementwise import Sub
|
||||
from mo.front.common.partial_infer.utils import int64_array
|
||||
from mo.front.common.replacement import FrontReplacementOp
|
||||
from mo.front.tf.graph_utils import create_op_with_const_inputs
|
||||
from mo.graph.graph import Graph, Node, rename_nodes
|
||||
|
||||
|
||||
class LogSoftmaxFrontReplacer(FrontReplacementOp):
|
||||
"""
|
||||
Replace LogSoftmax operation with ReduceMax + Sub + Exp + ReduceSum + Log + Sub.
|
||||
|
||||
More precisely, this transformation implements the following formulas of the calculation of LogSoftmax:
|
||||
|
||||
shifted_data = input_data - ReduceMax(input_data, axis), (1)
|
||||
output = shifted_data - Log(ReduceSum(Exp(shifted_data), axis)).
|
||||
|
||||
These formulas is used to calculate LogSoftmax in implementation of TensorFlow (see
|
||||
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/softmax_op_functor.h),
|
||||
Kaldi (see https://github.com/kaldi-asr/kaldi/blob/master/src/cudamatrix/cu-kernels.cu),
|
||||
MxNet (see https://github.com/apache/incubator-mxnet/blob/master/src/operator/nn/softmax-inl.h).
|
||||
|
||||
ONNX implements LogSoftmax according to formulas
|
||||
|
||||
flatten_data = Flatten(input_data, axis), (1')
|
||||
shifted_data = flatten_data - ReduceMax(flatten_data, 1),
|
||||
z = shifted_data - Log(ReduceSum(Exp(shifted_data), 1)),
|
||||
output = Reshape(z, input_data.shape)
|
||||
|
||||
(see https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/codegen/mti/math/logsoftmax.cc,
|
||||
https://github.com/microsoft/onnxruntime-tvm/blob/master/topi/include/topi/nn/softmax.h)
|
||||
|
||||
Formally speaking, the formula (1) is equivalent to the formula
|
||||
output = Log(SoftMax(input_data, axis)) (2)
|
||||
|
||||
But LogSoftMax is calculated according to formula (1) for better numeric stability.
|
||||
"""
|
||||
op = "LogSoftmax"
|
||||
enabled = True
|
||||
|
||||
def replace_op(self, graph: Graph, node: Node):
|
||||
node_name = node.soft_get('name', node.id)
|
||||
assert node.has_valid('axis'), 'The node "{}" does not have mandatory attribute "axis"'.format(node_name)
|
||||
|
||||
# Creating of ReduceMax -> Sub -> Exp block
|
||||
first_sub_node = Sub(graph, {'name': node_name + '/Sub_/first_'}).create_node()
|
||||
reduce_max_node = create_op_with_const_inputs(graph,
|
||||
ReduceMax,
|
||||
{1: int64_array([node.axis])},
|
||||
op_attrs={'name': node_name + '/ReduceMax_', 'keep_dims': True})
|
||||
reduce_max_node.out_port(0).connect(first_sub_node.in_port(1))
|
||||
|
||||
# Creating of Exp -> ReduceSum -> Log block
|
||||
exp_node = Exp(graph, {'name': node_name + '/Exp_'}).create_node()
|
||||
reduce_sum_node = create_op_with_const_inputs(graph,
|
||||
ReduceSum,
|
||||
{1: int64_array([node.axis])},
|
||||
op_attrs={'name': node_name + '/ReduceSum_', 'keep_dims': True})
|
||||
log_node = Log(graph, {'name': node_name + '/Log_'}).create_node()
|
||||
|
||||
first_sub_node.out_port(0).connect(exp_node.in_port(0))
|
||||
exp_node.out_port(0).connect(reduce_sum_node.in_port(0))
|
||||
reduce_sum_node.out_port(0).connect(log_node.in_port(0))
|
||||
|
||||
# Creating of the last Sub node
|
||||
second_sub_node = Sub(graph, {}).create_node()
|
||||
rename_nodes([(node, node_name + '/delete'), (second_sub_node, node_name)])
|
||||
log_node.out_port(0).connect(second_sub_node.in_port(1))
|
||||
first_sub_node.out_port(0).connect(second_sub_node.in_port(0))
|
||||
|
||||
# Correcting of input edges
|
||||
source = node.in_port(0).get_source()
|
||||
first_sub_node.in_port(0).connect(source)
|
||||
reduce_max_node.in_port(0).connect(source)
|
||||
|
||||
return [second_sub_node.id]
|
||||
@@ -1,86 +0,0 @@
|
||||
"""
|
||||
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
|
||||
|
||||
from generator import generator, generate
|
||||
|
||||
from extensions.front.LogSoftmax import LogSoftmaxFrontReplacer
|
||||
from mo.front.common.partial_infer.utils import int64_array
|
||||
from mo.utils.ir_engine.compare_graphs import compare_graphs
|
||||
from mo.utils.unittest.graph import build_graph
|
||||
|
||||
graph_node_attributes = {
|
||||
'placeholder': {'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
|
||||
'logsoftmax': {'type': None, 'kind': 'op', 'op': 'LogSoftmax', 'axis': -1},
|
||||
'output': {'kind': 'op', 'type': 'Result', 'op': 'Result'},
|
||||
}
|
||||
|
||||
|
||||
graph_edges = [
|
||||
('placeholder', 'logsoftmax'),
|
||||
('logsoftmax', 'output'),
|
||||
]
|
||||
|
||||
|
||||
graph_ref_node_attributes = {
|
||||
'placeholder': {'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
|
||||
'exp': {'type': 'Exp', 'kind': 'op', 'op': 'Exp'},
|
||||
'reduce_sum': {'type': 'ReduceSum', 'kind': 'op', 'op': 'ReduceSum', 'keep_dims': True},
|
||||
'reduce_max': {'type': 'ReduceMax', 'kind': 'op', 'op': 'ReduceMax', 'keep_dims': True},
|
||||
'log': {'type': 'Log', 'kind': 'op', 'op': 'Log'},
|
||||
'second_sub': {'type': 'Subtract', 'kind': 'op', 'op': 'Sub'},
|
||||
'reduce_sum_axis': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': None, 'shape': int64_array([1])},
|
||||
'reduce_max_axis': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': None, 'shape': int64_array([1])},
|
||||
'first_sub': {'type': 'Subtract', 'kind': 'op', 'op': 'Sub'},
|
||||
'output': {'kind': 'op', 'type': 'Result', 'op': 'Result'},
|
||||
}
|
||||
|
||||
|
||||
graph_ref_edges = [
|
||||
('placeholder', 'reduce_max', {'in': 0, 'out': 0}),
|
||||
('placeholder', 'first_sub', {'in': 0, 'out': 0}),
|
||||
('reduce_max', 'first_sub', {'in': 1}),
|
||||
('reduce_max_axis', 'reduce_max', {'in': 1}),
|
||||
('first_sub', 'exp', {'in': 0, 'out': 0}),
|
||||
('first_sub', 'second_sub', {'in': 0, 'out': 0}),
|
||||
('exp', 'reduce_sum', {'in': 0}),
|
||||
('reduce_sum_axis', 'reduce_sum', {'in': 1}),
|
||||
('reduce_sum', 'log'),
|
||||
('log', 'second_sub', {'in': 1}),
|
||||
('second_sub', 'output'),
|
||||
]
|
||||
|
||||
|
||||
@generator
|
||||
class LogSoftmaxReplacerTest(unittest.TestCase):
|
||||
@generate(*[(-1, 'NCHW'), (-1, 'NHWC'), (0, 'NHWC'),
|
||||
(0, 'NCHW'), (2, 'NCHW'), (2, 'NHWC'),
|
||||
(-2, 'NHWC'), (-2, 'NCHW')])
|
||||
def test_logsoftmax_replacer(self, axis, layout):
|
||||
graph = build_graph(nodes_attrs=graph_node_attributes, edges=graph_edges)
|
||||
graph_ref = build_graph(nodes_attrs=graph_ref_node_attributes,
|
||||
edges=graph_ref_edges,
|
||||
update_attributes={
|
||||
'reduce_max_axis': {'value': int64_array([axis])},
|
||||
'reduce_sum_axis': {'value': int64_array([axis])},
|
||||
})
|
||||
graph.graph['layout'] = layout
|
||||
graph.stage = 'front'
|
||||
LogSoftmaxFrontReplacer().find_and_replace_pattern(graph)
|
||||
(flag, resp) = compare_graphs(graph, graph_ref, 'output')
|
||||
self.assertTrue(flag, resp)
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from mo.ops.softmax import LogSoftmax
|
||||
from mo.ops.log_softmax import LogSoftmax
|
||||
from mo.front.extractor import FrontExtractorOp
|
||||
|
||||
|
||||
|
||||
@@ -36,10 +36,6 @@ class FlattenONNXToReshape(FrontReplacementSubgraph):
|
||||
"""
|
||||
enabled = True
|
||||
|
||||
def run_before(self):
|
||||
from extensions.front.LogSoftmax import LogSoftmaxFrontReplacer
|
||||
return [LogSoftmaxFrontReplacer]
|
||||
|
||||
def pattern(self):
|
||||
return dict(nodes=[('flatten', dict(op='FlattenONNX'))],
|
||||
edges=[])
|
||||
|
||||
@@ -18,7 +18,7 @@ from mo.graph.graph import Graph, Node, rename_nodes
|
||||
from mo.ops.flatten import FlattenONNX
|
||||
from mo.ops.reshape import Reshape
|
||||
from mo.ops.shape import Shape
|
||||
from mo.ops.softmax import LogSoftmax
|
||||
from mo.ops.log_softmax import LogSoftmax
|
||||
|
||||
|
||||
class LogSoftmaxONNXFrontReplacer(FrontReplacementOp):
|
||||
|
||||
@@ -16,7 +16,8 @@
|
||||
|
||||
from mo.front.extractor import FrontExtractorOp
|
||||
from mo.front.onnx.extractors.utils import onnx_attr
|
||||
from mo.ops.softmax import LogSoftmaxONNX, SoftmaxONNX
|
||||
from mo.ops.softmax import SoftmaxONNX
|
||||
from mo.ops.log_softmax import LogSoftmaxONNX
|
||||
|
||||
|
||||
class SoftmaxExtractor(FrontExtractorOp):
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
from mo.front.extractor import FrontExtractorOp
|
||||
from mo.ops.log_softmax import LogSoftmax
|
||||
|
||||
|
||||
class LogSoftmaxExtractor(FrontExtractorOp):
|
||||
op = 'LogSoftmax'
|
||||
enabled = True
|
||||
|
||||
@classmethod
|
||||
def extract(cls, node):
|
||||
# the default value for the TF LogSoftmax is -1
|
||||
axis = -1
|
||||
if 'axis' in node.pb.attr:
|
||||
axis = node.pb.attr['axis'].i
|
||||
LogSoftmax.update_node_stat(node, {'axis': axis})
|
||||
return cls.enabled
|
||||
@@ -15,7 +15,7 @@
|
||||
"""
|
||||
|
||||
from mo.front.extractor import FrontExtractorOp
|
||||
from mo.ops.softmax import LogSoftmax, Softmax
|
||||
from mo.ops.softmax import Softmax
|
||||
|
||||
|
||||
class SoftmaxExtractor(FrontExtractorOp):
|
||||
@@ -30,17 +30,3 @@ class SoftmaxExtractor(FrontExtractorOp):
|
||||
axis = node.pb.attr['axis'].i
|
||||
Softmax.update_node_stat(node, {'axis': axis})
|
||||
return cls.enabled
|
||||
|
||||
|
||||
class LogSoftmaxExtractor(FrontExtractorOp):
|
||||
op = 'LogSoftmax'
|
||||
enabled = True
|
||||
|
||||
@classmethod
|
||||
def extract(cls, node):
|
||||
# the default value for the TF LogSoftmax is -1
|
||||
axis = -1
|
||||
if 'axis' in node.pb.attr:
|
||||
axis = node.pb.attr['axis'].i
|
||||
LogSoftmax.update_node_stat(node, {'axis': axis})
|
||||
return cls.enabled
|
||||
Reference in new issue
Block a user