openvino/model-optimizer/extensions/middle/InterpolateSequenceToInterpolate.py

# Copyright (C) 2018-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import logging as log
import numpy as np
from typing import List

from extensions.ops.interpolate import Interpolate
from mo.front.common.partial_infer.utils import int64_array, shape_array
from mo.front.tf.graph_utils import create_op_with_const_inputs
from mo.graph.graph import Graph, Node, rename_nodes
from mo.middle.replacement import MiddleReplacementPattern
from mo.utils.error import Error
from mo.utils.utils import group_by_with_binary_predicate


def node_has_one_consumer(node: Node) -> bool:
    return len(node.out_port(0).get_destinations()) == 1


def is_next(first: Node, second: Node) -> bool:
    """
    This function checks if 'first' is predecessor of 'second'. The node 'first' is called to be
    a predecessor of the node 'second', if an output of 'first' is an input of 'second', and
    number of destinations of 'first' is equal to 1.
    :param first: an Interpolate layer
    :param second: another Interpolate layer
    :return: True, if 'first' is an predecessor of 'second', and False otherwise.
    """
    dests = first.out_port(0).get_destinations()
    if node_has_one_consumer(first):
        return second.id == dests[0].node.id
    elif first.soft_get('maybe_part_of_sequence', False):
        return len(dests) == 2 and second.id in [d.node.id for d in dests]
    return False


class CanBeFused:
    def __init__(self):
        # We need to accumulate set of axes of compared nodes, because there can be a sequence of a set of axes
        #   {i}{j}{i}
        self.accumulated_axes = set()
        self.default_values_for_opset4 = {
            'mode': None,
            'shape_calculation_mode': None,
            'coordinate_transformation_mode': 'half_pixel',
            'nearest_mode': 'round_prefer_floor',
            'antialias': 0,
            'cube_coeff': -0.75
        }
        self.default_pads = int64_array([0])

    def _compare_attributes_of_interpolate1(self, first: Node, second: Node) -> bool:
        """
        This function checks whether attributes of Interpolate-1 nodes first and second are identical
        (except attribute 'axes').
        :param first: the first of compared nodes
        :param second: the second of compared nodes
        :return: True, if attributes of nodes are identical and False otherwise
        """
        # If some of attributes 'mode', 'align_corners', 'antialias', 'pads_begin', 'pads_end' are different,
        # then attributes of nodes are not identical.
        op = Interpolate(graph=first.graph, attrs={})
        for attr in ['mode', 'align_corners', 'antialias', 'pads_begin', 'pads_end']:
            if first.soft_get(attr, default=op.attrs[attr]) != second.soft_get(attr, default=op.attrs[attr]):
                return False
        return True

    def _compare_attributes_of_interpolate4(self, first: Node, second: Node) -> bool:
        """
        This function checks whether attributes of Interpolate-4 nodes first and second are identical.
        :param first: the first of compared nodes
        :param second: the second of compared nodes
        :return: True, if attributes of nodes are identical and False otherwise
        """
        # If some of attributes 'mode', 'coordinate_transformation_mode', 'nearest_mode', 'antialias', 'cube_coeff'
        # are different, then attributes of first and second are not identical.
        for attr in self.default_values_for_opset4.keys():
            default_value = self.default_values_for_opset4[attr]
            if first.soft_get(attr, default=default_value) != second.soft_get(attr, default=default_value):
                return False

        # If attributes 'pads_begin' or 'pads_end' of nodes first and second are different, then attributes
        # of first and second are not identical.
        for attr in ['pads_begin', 'pads_end']:
            if not np.array_equal(first.soft_get(attr, default=self.default_pads),
                                  second.soft_get(attr, default=self.default_pads)):
                return False
        return True

    def _compare_attributes(self, first: Node, second: Node) -> bool:
        """
        This function checks whether attributes of nodes first and second are identical (except attribute 'axes').
        :param first: the first of compared nodes
        :param second: the second of compared nodes
        :return: True, if attributes of nodes are identical and False otherwise
        """
        # If opsets of nodes are different, then nodes have different attributes.
        fst_opset = first.get_opset()
        snd_opset = second.get_opset()
        if fst_opset != snd_opset:
            return False

        if fst_opset not in ['opset1', 'opset4']:
            fst_name = first.soft_get('name', first.id)
            snd_name = second.soft_get('name', second.id)
            raise Error('Unsupported opset {} for nodes with names {} and {}'.format(fst_opset, fst_name, snd_name))

        if fst_opset == 'opset1':
            return self._compare_attributes_of_interpolate1(first, second)
        else:
            return self._compare_attributes_of_interpolate4(first, second)

    def __call__(self, first: Node, second: Node) -> bool:
        """
        This function checks whether Interpolate nodes 'first' and 'second' can be fused.
        :param first: the first of fused nodes
        :param second: the second of fused nodes
        :return: True, if nodes can be fused, and False otherwise
        """
        if not (is_next(first, second) and self._compare_attributes(first, second)):
            self.accumulated_axes = set()
            return False

        fst_axes = set([a for a in Interpolate.get_axes(first)])
        snd_axes = set([a for a in Interpolate.get_axes(second)])

        self.accumulated_axes = self.accumulated_axes | fst_axes

        # If the set of accumulated axes and the set of axes of 'second' do not intersect then nodes can be fused,
        # because interpolations with respect to various axes do not affect each other.
        if not(self.accumulated_axes & snd_axes):
            return True

        # Otherwise, nodes cannot be fused.
        self.accumulated_axes = set()
        return False


def get_interpolate_attributes(node: Node) -> dict:
    opset_to_default_values = {
        'opset1': {
            'mode': None,
            'align_corners': 0,
            'antialias': 0,
            'pads_begin': 0,
            'pads_end': 0,
            'version': 'opset1'
        },
        'opset4': {
            'mode': None,
            'shape_calculation_mode': None,
            'antialias': 0,
            'pads_begin': int64_array([0]),
            'pads_end': int64_array([0]),
            'coordinate_transformation_mode': 'half_pixel',
            'nearest_mode': 'round_prefer_floor',
            'cube_coeff': -0.75,
            'version': 'opset4'
        },
    }
    opset = node.get_opset()
    result = {}
    if opset in opset_to_default_values:
        default_values = opset_to_default_values[opset]
        for attr in default_values.keys():
            value = node.soft_get(attr, default=default_values[attr])
            result[attr] = value
        return result
    else:
        raise Error('Unsupported opset {} for node with name {}.'.format(opset, node.soft_get('name', node.id)))


def replace_sequence(seq: List[Node], graph: Graph):
    """
    This function replaces a sequence of consecutive Interpolate layers with one Interpolate layer,
    if modes of all nodes of a sequence are the same.
    :param seq: sequence of Interpolate layers
    :param graph: graph to which nodes of seq belong
    :return: Nothing
    """
    if not seq:
        return
    if len(seq) == 1:
        return

    modes = set([n.mode for n in seq])
    if len(modes) != 1:
        return

    dims_and_scales_ = []
    # Each element of the list dims_and_scales_ is a pair
    #      (axis, output size for this axis) (opset1)
    # or
    #      (axis, output size for this axis, output scales for this axis) (opset4)
    if seq[0].get_opset() == 'opset1':
        for interp in seq:
            dims_and_scales_.extend(zip(Interpolate.get_axes(interp),
                                        interp.in_port(1).get_connection().get_source().data.get_value()))

        axis_to_size = sorted(list(dict(dims_and_scales_).items()), key=lambda x: x[0])
        axes_of_node = int64_array([z[0] for z in axis_to_size])
        sizes = shape_array([z[1] for z in axis_to_size])
        scales = np.ones(len(axis_to_size))
    else:
        for interp in seq:
            dims_and_scales_.extend(zip(Interpolate.get_axes(interp),
                                        interp.in_port(1).get_connection().get_source().data.get_value(),
                                        interp.in_port(2).get_connection().get_source().data.get_value()))

        axis_to_size = sorted(dims_and_scales_, key=lambda x: x[0])
        axes_of_node = int64_array([z[0] for z in axis_to_size])
        sizes = shape_array([z[1] for z in axis_to_size])
        scales = np.array([z[2] for z in axis_to_size])

    fst_interp_node = seq[0]
    last_interp_node = seq[-1]
    last_interp_node_name = last_interp_node.soft_get('name', last_interp_node.id)
    attributes = get_interpolate_attributes(fst_interp_node)

    opset = fst_interp_node.get_opset()
    if opset == 'opset1':
        attributes['axes'] = axes_of_node
        interp_node = create_op_with_const_inputs(graph, Interpolate, {1: sizes}, attributes)

        fst_interp_connection = fst_interp_node.in_port(0).get_connection()
        fst_interp_connection.set_destination(interp_node.in_port(0))

        last_interp_node.out_port(0).get_connection().set_source(interp_node.out_port(0))
    else:
        attributes['in_ports_count'] = 4
        interp_node = create_op_with_const_inputs(graph, Interpolate,
                                                  {1: sizes, 2: scales, 3: axes_of_node},
                                                  attributes)

        fst_interp_connection = fst_interp_node.in_port(0).get_connection()
        fst_interp_connection.set_destination(interp_node.in_port(0))

        last_interp_node.out_port(0).get_connection().set_source(interp_node.out_port(0))

    rename_nodes([(last_interp_node, last_interp_node_name + '/delete'), (interp_node, last_interp_node_name)])


class InterpolateSequenceToInterpolate(MiddleReplacementPattern):
    """
    This transformation replaces a sequence of Interpolate layers by one Interpolate layer.
    """
    enabled = True

    def run_before(self):
        from extensions.middle.UpsampleToResample import UpsampleToResample
        return [UpsampleToResample]

    def find_and_replace_pattern(self, graph: Graph):
        log.debug('Enabled replacement of a sequence of Interpolate layers with one Interpolate layer.')
        interps = [n for n in graph.pseudo_topological_sort() if n.kind == 'op' and n.op == 'Interpolate']
        fuser = CanBeFused()
        sequences = group_by_with_binary_predicate(interps, fuser)
        for seq in sequences:
            replace_sequence(seq, graph)