model-optimizer/extensions/middle/GatherNdNormalizer.py

   1 """
   2  Copyright (c) 2019 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16 import logging as log
  17
  18 import numpy as np
  19
  20 from extensions.ops.gather import Gather
  21 from mo.front.common.partial_infer.utils import int64_array
  22 from mo.graph.graph import Graph
  23 from mo.middle.replacement import MiddleReplacementPattern
  24 from mo.ops.const import Const
  25 from mo.ops.reshape import Reshape
  26
  27
  28 class GatherNdNormalize(MiddleReplacementPattern):
  29     """
  30     Hot fix for new speech-to-text model enabling while GatherND is not implemented in IE.
  31     We can replace GatherNd to Reshape + Gather in case when GatherNd indices have just one
  32     meaningful dimension.
  33     """
  34     enabled = True
  35     force_clean_up = True
  36
  37     def run_before(self):
  38         from extensions.middle.BlockLSTMtoLSTMSequence import BlockLSTMtoLSTMSequence
  39         return [BlockLSTMtoLSTMSequence]
  40
  41     def run_after(self):
  42         from extensions.middle.pass_separator import MiddleStart
  43         return [MiddleStart]
  44
  45     def pattern(self):
  46         return dict(
  47             nodes=[('GatherNd', dict(kind='op', op='GatherNd'))],
  48             edges=[]
  49         )
  50
  51     @staticmethod
  52     def indices_check(indices: np.array, input_shape: tuple):
  53         """
  54         Check that indices have just one meaningful dimension and all other dimensions of input have size 1.
  55         """
  56         n_dims = indices.shape[-1]
  57         non_zero = None
  58         for i in range(n_dims):
  59             if not all(np.take(indices, indices=[i], axis=-1) == 0):
  60                 if non_zero is None:
  61                     non_zero = i
  62                 else:
  63                     return None
  64             else:
  65                 if input_shape[i] != 1:
  66                     return None
  67         return non_zero
  68
  69     def replace_pattern(self, graph: Graph, match: dict):
  70         gather = match['GatherNd']
  71         input_shape = gather.in_node(0).shape
  72         indices = gather.in_node(1).value
  73         if indices is None:
  74             # We can't do such special pass without indices value
  75             return
  76
  77         # 0. All needed checks that we can replace GatherNd by Gather
  78         gather_idx = self.indices_check(indices, input_shape)
  79         if gather_idx is None:
  80             log.warning('Node {} with op=GatherNd  can\'t be normalized to op=Gather.'.format(gather.name))
  81             return
  82
  83         # 1. Add Reshape and connect
  84         new_shape = int64_array([-1] + list(input_shape[indices.shape[-1]:]))
  85         reshape = Reshape(graph, {'name': gather.name + '/Reshape_for_GatherNd/', 'dim': new_shape, }).create_node()
  86         gather.in_port(0).get_connection().set_destination(reshape.in_port(0))
  87
  88         # 2. Change indices from Nd to 1d:
  89         new_indices = np.reshape(np.take(indices, indices=[gather_idx], axis=-1), [-1])
  90         new_indices_const = Const(graph, dict(value=new_indices)).create_node()
  91
  92         # 3. Create new Gather operation and reconnect all inputs/outputs
  93         new_gather = Gather(graph, {'name': gather.name + '/NewGather/', 'axis': 0}).create_node()
  94         reshape.out_port(0).connect(new_gather.in_port(0))
  95         new_indices_const.out_port(0).connect(new_gather.in_port(1))
  96
  97         gather.out_port(0).get_connection().set_source(new_gather.out_port(0))
  98
  99         # 4. Remove old Gather node
 100         graph.remove_node(gather.id)