211 lines
9.1 KiB
Python
211 lines
9.1 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from extensions.ops.elementwise import Add
|
|
from extensions.ops.gather import Gather
|
|
from extensions.ops.range import Range
|
|
from mo.front.common.partial_infer.utils import int64_array
|
|
from mo.front.tf.graph_utils import create_op_node_with_second_input
|
|
from mo.graph.graph import Node
|
|
from mo.graph.port import Port
|
|
from mo.ops.concat import Concat
|
|
from mo.ops.const import Const
|
|
from mo.ops.shape import Shape
|
|
from mo.ops.squeeze import Squeeze
|
|
|
|
|
|
def get_canonical_axis_index_node(rank: Node, axis: int) -> Node:
|
|
"""
|
|
Returns positive axis value
|
|
|
|
:param rank: the node of 0D output shape to get rank of tensor from
|
|
:param axis: integer value from [-rank; rank - 1]
|
|
:return: node producing positive integer value of axis
|
|
"""
|
|
graph = rank.graph
|
|
name = rank.soft_get('name', rank.id)
|
|
if axis < 0:
|
|
axis = Const(graph, {'name': name + '/negative_axis', 'value': int64_array(axis)}).create_node()
|
|
add = Add(graph, {'name': name + '/positive_axis'}).create_node()
|
|
rank.out_port(0).connect(add.in_port(0))
|
|
axis.out_port(0).connect(add.in_port(1))
|
|
return add
|
|
else:
|
|
return Const(graph, {'name': name + '/positive_axis', 'value': int64_array(axis)}).create_node()
|
|
|
|
|
|
def get_range_node_of_idxs(rank: Node, begin: int, end: int,
|
|
include_begin: bool = True, include_end: bool = False) -> Node:
|
|
"""
|
|
Returns node that produces 1D output of values of range from begin to end (ex)/(in)cluding begin or end point
|
|
|
|
:param rank: the node of 0D output shape to get rank of tensor from
|
|
:param begin: integer value from [-rank; rank - 1]
|
|
:param end: integer value from [-rank; +rank]
|
|
:param include_begin: boolean flag to include or exclude start point from range output
|
|
:param include_end: boolean flag to include or exclude end point from range output
|
|
:return: range node producing 1D output
|
|
"""
|
|
graph = rank.graph
|
|
name = rank.soft_get('name', rank.id)
|
|
|
|
start_idx = get_canonical_axis_index_node(rank, begin)
|
|
end_idx = get_canonical_axis_index_node(rank, end)
|
|
|
|
if not include_begin:
|
|
const = Const(graph, {'value': int64_array(1), 'name': name + '/exclude_begin/value'}).create_node()
|
|
add = Add(graph, {'name': name + '/exclude_begin'}).create_node()
|
|
start_idx.out_port(0).connect(add.in_port(0))
|
|
const.out_port(0).connect(add.in_port(1))
|
|
start_idx = add
|
|
|
|
if include_end:
|
|
const = Const(graph, {'value': int64_array(1), 'name': name + '/including_end/value'}).create_node()
|
|
add = Add(graph, {'name': name + '/including_end'}).create_node()
|
|
end_idx.out_port(0).connect(add.in_port(0))
|
|
const.out_port(0).connect(add.in_port(1))
|
|
end_idx = add
|
|
|
|
delta = Const(graph, {'name': name + '/delta', 'value': int64_array(1)}).create_node()
|
|
range_node = Range(graph, {'name': name + '/range_idxs'}).create_node()
|
|
|
|
start_idx.out_port(0).connect(range_node.in_port(0))
|
|
end_idx.out_port(0).connect(range_node.in_port(1))
|
|
delta.out_port(0).connect(range_node.in_port(2))
|
|
|
|
return range_node
|
|
|
|
|
|
def get_shape_values_by_indices_node(shape_node: Node, indices_node: Node) -> Node:
|
|
"""
|
|
The function returns a node that produces values of the specified indices node of the input node 'shape_node'
|
|
|
|
:param shape_node: the node of 1D output shape to get elements from
|
|
:param indices_node: the node of 1D output shape with the list of element indices to get
|
|
:return: node producing required elements of the node
|
|
"""
|
|
graph = shape_node.graph
|
|
axis = Const(graph, {'value': int64_array(0), 'name': shape_node.name + '/Axis'}).create_node()
|
|
gather_node = Gather(graph, {'name': shape_node.name + '/Gather'}).create_node()
|
|
|
|
shape_node.out_port(0).connect(gather_node.in_port(0))
|
|
indices_node.out_port(0).connect(gather_node.in_port(1))
|
|
axis.out_port(0).connect(gather_node.in_port(2))
|
|
return gather_node
|
|
|
|
|
|
def node_to_get_shape_value_of_indices(shape_node: Node, indices: list) -> Node:
|
|
"""
|
|
The function returns a node that produces values of the specified indices of the input node 'shape_node'
|
|
|
|
:param shape_node: the node of 1D output shape to get elements from
|
|
:param indices: the list of element indices to get
|
|
:return: node producing required elements of the node
|
|
"""
|
|
graph = shape_node.graph
|
|
indices_node = Const(graph, {'value': int64_array(indices), 'name': shape_node.name + '/Indices'}).create_node()
|
|
|
|
gather_node = get_shape_values_by_indices_node(shape_node, indices_node)
|
|
return gather_node
|
|
|
|
|
|
def get_shape_values_by_range_idxs(shape: Node, rank: Node, begin: int, end: int,
|
|
include_begin: bool = True, include_end: bool = False):
|
|
"""
|
|
Gathers shape values that are represented by range from begin to end (in)/(ex)cluding begin or end point
|
|
|
|
:param shape: the node of 1D output shape to get elements from
|
|
:param rank: the node of 0D output shape to get rank of tensor from
|
|
:param begin: integer value from [-rank; rank - 1]
|
|
:param end: integer value from [-rank; +rank]
|
|
:param include_begin: boolean flag to include or exclude start point from range output
|
|
:param include_end: boolean flag to include or exclude end point from range output
|
|
:return: gather node producing 1D output
|
|
"""
|
|
range_node = get_range_node_of_idxs(rank, begin, end, include_begin=include_begin, include_end=include_end)
|
|
return get_shape_values_by_indices_node(shape, range_node)
|
|
|
|
|
|
def node_to_get_batch_value(shape_node: Node) -> Node:
|
|
"""
|
|
The function returns a node that produces the batch value which is usually the element of the shape with index 0
|
|
:param shape_node: the node of 1D output shape to get batch from
|
|
:return: the node producing batch value
|
|
"""
|
|
return node_to_get_shape_value_of_indices(shape_node, [0])
|
|
|
|
|
|
def node_to_get_features_dimension_value(shape_node: Node) -> Node:
|
|
"""
|
|
The function returns a node that produces the feature dimension value
|
|
:param shape_node: the node of 1D output shape to get the feature dimension value from
|
|
:return: the node producing feature dimension value
|
|
"""
|
|
layout = shape_node.graph.graph['layout']
|
|
if layout == 'NCHW':
|
|
return node_to_get_shape_value_of_indices(shape_node, [1])
|
|
elif layout == 'NHWC':
|
|
return node_to_get_shape_value_of_indices(shape_node, [-1])
|
|
else:
|
|
assert 'Unsupported layout "{}"'.format(layout)
|
|
|
|
|
|
def node_to_get_spatial_dimensions_value(shape_node: Node) -> Node:
|
|
"""
|
|
The function returns a node that produces the spatial dimension values
|
|
:param shape_node: the node of 1D output shape to get the spatial dimension values from
|
|
:return: the node producing the spatial dimension values
|
|
"""
|
|
layout = shape_node.graph.graph['layout']
|
|
shape = shape_node.in_port(0).get_connection().get_source().data.get_shape()
|
|
assert shape is not None, 'The shape must be inferred before running this function'
|
|
|
|
if layout == 'NCHW':
|
|
return node_to_get_shape_value_of_indices(shape_node, list(range(2, len(shape))))
|
|
elif layout == 'NHWC':
|
|
return node_to_get_shape_value_of_indices(shape_node, list(range(1, len(shape) - 1)))
|
|
else:
|
|
assert 'Unsupported layout "{}"'.format(layout)
|
|
|
|
|
|
def new_shape_node_from_shape_nodes(input_shape_nodes: list):
|
|
"""
|
|
The function returns a node producing 1D tensor with concatenated shapes produced by nodes from "input_shape_nodes"
|
|
:param input_shape_nodes: list of nodes producing 1D tensors
|
|
:return: the node producing concatenated values of nodes from the "input_shape_nodes"
|
|
"""
|
|
assert len(input_shape_nodes) > 0, 'The list of input shape nodes should be non-empty'
|
|
new_shape_node = Concat(input_shape_nodes[0].graph,
|
|
{'axis': 0,
|
|
'name': input_shape_nodes[0].soft_get('name', input_shape_nodes[0].id) + '/shapes_concat'}
|
|
).create_node()
|
|
|
|
for ind, input_node in enumerate(input_shape_nodes):
|
|
new_shape_node.add_input_port(ind)
|
|
new_shape_node.in_port(ind).connect(input_node.out_port(0))
|
|
return new_shape_node
|
|
|
|
|
|
def get_shape_and_rank_nodes_by_port(port: Port, return_as_a_scalar: bool = True):
|
|
"""
|
|
The function returns nodes producing shape and rank of the data from the desired port in order to use those
|
|
operations on the middle/back phase
|
|
:param port: Port object that specifies node output port
|
|
:param return_as_a_scalar: boolean flag to return 1D or 0D rank
|
|
:return: shape and rank nodes
|
|
"""
|
|
input_node_name = port.node.soft_get('name', port.node.id)
|
|
graph = port.node.graph
|
|
|
|
shape = Shape(graph, dict(name=input_node_name + '/ShapeOf')).create_node()
|
|
rank_1_d = Shape(graph, dict(name=input_node_name + '/1dRankOf')).create_node()
|
|
rank_1_d.in_port(0).connect(shape.out_port(0))
|
|
shape.in_port(0).connect(port)
|
|
if not return_as_a_scalar:
|
|
return shape, rank_1_d
|
|
|
|
rank = create_op_node_with_second_input(graph, Squeeze, int64_array([0]), {'name': input_node_name + '/0dRankOf'},
|
|
rank_1_d)
|
|
return shape, rank
|
|
|