67 lines
3.6 KiB
Python
67 lines
3.6 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
|
|
from extensions.ops.Cast import Cast
|
|
from extensions.ops.elementwise import Mul, Sub
|
|
from mo.front.common.partial_infer.utils import int64_array
|
|
from mo.front.tf.graph_utils import create_op_with_const_inputs
|
|
from mo.graph.graph import Graph, rename_nodes
|
|
from mo.middle.passes.convert_data_type import data_type_str_to_np
|
|
from mo.middle.replacement import MiddleReplacementPattern
|
|
from mo.ops.reshape import Reshape
|
|
|
|
|
|
class DequantizeLinearResolver(MiddleReplacementPattern):
|
|
"""
|
|
Transformation result depends on the axis value.
|
|
If the axis is not set or x_scale input is scalar or 1D tensor with one element then DequantizeLinear is
|
|
replaced with the sub-graph which can be expressed with the following formula:
|
|
y = (x - x_zero_point) * x_scale
|
|
In other cases DequantizeLinear can be replace to formula with addition reshape x_zero_point and x_scale.
|
|
Target shape for reshape depend on axis.
|
|
"""
|
|
enabled = True
|
|
graph_condition = [lambda graph: graph.graph['layout'] == 'NCHW']
|
|
|
|
def find_and_replace_pattern(self, graph: Graph):
|
|
for dequantize_node in graph.get_op_nodes(op='DequantizeLinear'):
|
|
node_name = dequantize_node.soft_get('name', dequantize_node.id)
|
|
axis = dequantize_node.soft_get('axis', None)
|
|
scale_y_shape = dequantize_node.in_port(1).data.get_shape()
|
|
model_data_type = data_type_str_to_np(graph.graph['cmd_params'].data_type)
|
|
cast = Cast(graph, {'dst_type': model_data_type, 'name': node_name + '/Cast'}).create_node()
|
|
dequantize_node.in_port(0).get_connection().set_destination(cast.in_port(0))
|
|
mul = Mul(graph, {}).create_node()
|
|
|
|
is_second_port_connected = dequantize_node.is_in_port_connected(2)
|
|
if is_second_port_connected:
|
|
sub = Sub(graph, {'name': node_name + '/Sub'}).create_node()
|
|
cast.out_port(0).connect(sub.in_port(0))
|
|
dequantize_node.in_port(2).get_connection().set_destination(sub.in_port(1))
|
|
sub.out_port(0).connect(mul.in_port(0))
|
|
else:
|
|
cast.out_port(0).connect(mul.in_port(0))
|
|
|
|
dequantize_node.in_port(1).get_connection().set_destination(mul.in_port(1))
|
|
dequantize_node.out_port(0).get_connection().set_source(mul.out_port(0))
|
|
rename_nodes([(dequantize_node, node_name + '/TBD'), (mul, node_name)])
|
|
|
|
assert scale_y_shape is not None
|
|
if axis is not None and len(scale_y_shape) > 0 and scale_y_shape[0] > 1:
|
|
input_shape = cast.in_port(0).data.get_shape()
|
|
target_shape = np.ones(len(input_shape), np.int64)
|
|
target_shape[axis] = input_shape[axis]
|
|
|
|
mul_reshape = create_op_with_const_inputs(graph, Reshape, {1: int64_array(target_shape)},
|
|
{'name': node_name + '/Reshape/Mul'})
|
|
mul.in_port(1).get_connection().set_destination(mul_reshape.in_port(0))
|
|
mul_reshape.out_port(0).connect(mul.in_port(1))
|
|
|
|
if is_second_port_connected:
|
|
sub_reshape = create_op_with_const_inputs(graph, Reshape, {1: int64_array(target_shape)},
|
|
{'name': node_name + '/Reshape/Sub'})
|
|
sub.in_port(1).get_connection().set_destination(sub_reshape.in_port(0))
|
|
sub_reshape.out_port(0).connect(sub.in_port(1))
|