134 lines
5.3 KiB
Python
134 lines
5.3 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import logging as log
|
|
|
|
import numpy as np
|
|
|
|
from extensions.back.OptimizeTransposeReshapeSequence import set_reshape_new_output_shape
|
|
from mo.back.replacement import BackReplacementPattern
|
|
from mo.front.common.partial_infer.utils import int64_array
|
|
from mo.graph.graph import Graph
|
|
|
|
|
|
def sequential_dims(order: np.array):
|
|
"""
|
|
Returns first instance (counting from left) of the sequential dimensions in the 'order'
|
|
:param order: order to look for sequential dims
|
|
:return: list of indices of the sequential dimensions. If sequential dimensions are not found then return None.
|
|
"""
|
|
start_ind = cur_ind = 0
|
|
while cur_ind + 1 < len(order):
|
|
if order[cur_ind] + 1 == order[cur_ind + 1]:
|
|
cur_ind += 1
|
|
else:
|
|
if start_ind < cur_ind:
|
|
return list(range(start_ind, cur_ind + 1))
|
|
cur_ind += 1
|
|
start_ind = cur_ind
|
|
if start_ind < cur_ind:
|
|
return list(range(start_ind, cur_ind + 1))
|
|
return None
|
|
|
|
|
|
def merge_permute_order_dimensions(dims: list, permute_order: np.array):
|
|
"""
|
|
Creates updated permutation for a given permutation order and the *input* dimension indices to be merged into one.
|
|
:param dims: the input tensor dimensions indices to merge
|
|
:param permute_order: the permutation order
|
|
:return: the new permutation order after merging of the specified dimensions into one
|
|
"""
|
|
assert len(dims) >= 2
|
|
new_permute_order = list()
|
|
for permute_index in permute_order:
|
|
if permute_index < permute_order[dims[0]]:
|
|
new_permute_order.append(permute_index)
|
|
elif permute_index > permute_order[dims[-1]]:
|
|
new_permute_order.append(permute_index - len(dims) + 1)
|
|
elif permute_index == permute_order[dims[0]]:
|
|
new_permute_order.append(permute_order[dims[0]])
|
|
return int64_array(new_permute_order)
|
|
|
|
|
|
def merge_dims(dims_to_merge: np.array, shape: np.array):
|
|
"""
|
|
Merge several sequential specified dims into one.
|
|
|
|
The function does not support magic number "0" in the 'shape'.
|
|
:param dims_to_merge: the dimensions indices to merge
|
|
:param shape: shape to merge
|
|
:return: new shape with merged specified dims
|
|
"""
|
|
for ind in range(len(dims_to_merge) - 1):
|
|
assert dims_to_merge[ind] + 1 == dims_to_merge[ind + 1], 'The dims to merge must be sequential'
|
|
assert 0 not in shape, 'The value 0 is not supported during merging of the shape'
|
|
|
|
result = list()
|
|
if dims_to_merge[0] != 0:
|
|
result.extend(shape[:dims_to_merge[0]])
|
|
# handle magic number "-1"
|
|
if -1 in shape[dims_to_merge]:
|
|
result.append(-1)
|
|
else:
|
|
result.append(np.prod(shape[dims_to_merge]))
|
|
if dims_to_merge[-1] + 1 != len(shape):
|
|
result.extend(shape[dims_to_merge[-1] + 1:])
|
|
return int64_array(result)
|
|
|
|
|
|
class ReduceTransposeDimensions(BackReplacementPattern):
|
|
"""
|
|
Transformation looks for the Transpose layers with sequential dimensions in the permutation order and merges them into
|
|
one thus reducing the number of dimensions. The transformation is applied to 5D+ permutations only.
|
|
"""
|
|
enabled = False
|
|
|
|
def run_after(self):
|
|
from extensions.back.OptimizeTransposeReshapeSequence import OptimizeTransposeReshapeSequence
|
|
return [OptimizeTransposeReshapeSequence]
|
|
|
|
def run_before(self):
|
|
from extensions.back.ReshapeMutation import ReshapeMutation
|
|
return [ReshapeMutation]
|
|
|
|
@staticmethod
|
|
def pattern():
|
|
return dict(
|
|
nodes=[
|
|
('reshape_1', dict(kind='op', type='Reshape')), # TODO change to reshape-like
|
|
('reshape_1_data', dict(kind='data')),
|
|
('permute', dict(kind='op', type='Transpose')),
|
|
('permute_data', dict(kind='data')),
|
|
('reshape_2', dict(kind='op', type='Reshape')),
|
|
],
|
|
edges=[
|
|
('reshape_1', 'reshape_1_data'),
|
|
('reshape_1_data', 'permute'),
|
|
('permute', 'permute_data'),
|
|
('permute_data', 'reshape_2'),
|
|
]
|
|
)
|
|
|
|
@staticmethod
|
|
def replace_pattern(graph: Graph, match: dict):
|
|
permute_node = match['permute']
|
|
reshape_1_node = match['reshape_1']
|
|
order = permute_node.in_port(1).data.get_value().copy()
|
|
if len(order) >= 5:
|
|
log.debug('Trying to merge dimensions of the Transpose layer "{}"'.format(permute_node.soft_get('name')))
|
|
seq_dims = sequential_dims(order)
|
|
while seq_dims is not None:
|
|
permute_input_shape = permute_node.in_port(0).data.get_shape().copy()
|
|
# calculate new Transpose order and output of the Reshape layer
|
|
new_reshape_dims = merge_dims(order[seq_dims], permute_input_shape)
|
|
new_permute_order = merge_permute_order_dimensions(seq_dims, order)
|
|
|
|
assert reshape_1_node.has('dim')
|
|
# update data Transpose and Reshape attributes and data nodes shapes
|
|
set_reshape_new_output_shape(reshape_1_node, new_reshape_dims)
|
|
|
|
permute_node.in_port(1).data.set_value(new_permute_order)
|
|
permute_node.infer(permute_node)
|
|
order = permute_node.in_port(1).data.get_value().copy()
|
|
seq_dims = sequential_dims(order)
|