model-optimizer/extensions/middle/EltwiseInputNormalization.py

   1 """
   2  Copyright (c) 2018-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
  17 import networkx as nx
  18 import numpy as np
  19
  20 from extensions.middle.EltwiseInputReshape import EltwiseInputReshape
  21 from mo.graph.graph import Node, Graph
  22 from mo.middle.replacement import MiddleReplacementPattern
  23
  24
  25 class EltwiseInputNormalize(EltwiseInputReshape, MiddleReplacementPattern):
  26     # This pass should be called directly from pipeline before layout change and other permutations
  27     enabled = False
  28
  29     def find_and_replace_pattern(self, graph: Graph):
  30         eltwise_nodes = [Node(graph, node) for node in graph.node if Node(graph, node).soft_get('type') == 'Eltwise']
  31         # Iterating over all Eltwise operations and check that every input has similar shape
  32         # in case of different shapes, we inserts new_shape attribute and then call EltwiseInputReshape extension
  33         # that insert reshapes (in case of not constant nodes) or directly reshapes values in data nodes for specified
  34         # shape
  35         for node in eltwise_nodes:
  36             output_shape = node.out_node().shape
  37             for in_node in node.in_nodes().values():
  38                 if len(in_node.shape) != len(output_shape):
  39                     # Set edge attribute new_shape for further transformation pass
  40                     new_shape = in_node.shape
  41                     for x in range(len(output_shape) - len(in_node.shape)):
  42                         new_shape = np.insert(new_shape, 0, 1)
  43
  44                     nx.set_edge_attributes(G=node.graph,
  45                                            values={(in_node.id, node.id, 0): new_shape},
  46                                            name='new_shape')
  47
  48         super().find_and_replace_pattern(graph)