model-optimizer/mo/front/tf/graph_utils.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 collections
  18 import logging as log
  19
  20 import numpy as np
  21
  22 from mo.front.extractor import update_attrs
  23 from mo.graph.graph import Node, Graph
  24 from mo.ops.activation import Activation
  25 from mo.ops.concat import Concat
  26 from mo.ops.const import Const
  27 from mo.ops.convolution import Convolution
  28 from mo.ops.crop import Crop
  29 from mo.ops.reshape import Reshape
  30 from mo.ops.softmax import Softmax
  31 from mo.utils.error import Error
  32
  33
  34 def squeeze_reshape_and_concat(start_nodes: list):
  35     """
  36     The function looks for Reshape ops after the 'start_nodes' with 4D output and remove the dimension with index 2
  37     which should be equal to 1. This is a workaround to make tensor 3D so it's shape will not be transposed during the
  38     IR generation. The problem arises when bounding boxes predictions are reshaped from [1, 1, 1, X] to
  39     [1, X / 4, 1, 4]. The result tensor should not be transposed because after transpose it will have shape
  40     [1, 4, X / 4, 1] and the concatenation over dimension with index 2 will produce incorrect tensor.
  41     Also the function looks for Concat ops and change the concat dimension from 2 to 1.
  42     :param start_nodes: list of nodes to start search from.
  43     :return: None
  44     """
  45     q = collections.deque()
  46     q.extend(start_nodes)
  47     while len(q) != 0:
  48         cur_node = q.popleft()
  49         if cur_node.has_valid('type'):
  50             if cur_node.type == 'DetectionOutput':  # do not go beyond the DetectionOutput node
  51                 continue
  52             if cur_node.type == 'Reshape' and len(cur_node.out_node().shape) == 4:
  53                 log.debug("Found Reshape op with 4D output {}".format(cur_node.id))
  54                 if cur_node.in_node(1).has_valid('value') and cur_node.in_node(1).value is not None:
  55                     new_shape = cur_node.in_node(1).value
  56                     assert new_shape[2] == 1
  57                     new_shape = np.delete(new_shape, 2)
  58                     cur_node.in_node(1).value = new_shape
  59                     cur_node.in_node(1).shape = np.array(new_shape.shape, dtype=np.int64)
  60                     cur_node['dim'] = new_shape.copy()
  61                     # run infer function once again
  62                     cur_node.infer(cur_node)
  63                 else:
  64                     log.warning("The reshape size is not defined!")
  65             if cur_node.type == 'Concat' and len(cur_node.out_node().shape) == 4:
  66                 log.debug("Found Concat op with 4D output {}".format(cur_node.id))
  67                 cur_node.axis = 1
  68                 # run infer function once again
  69                 cur_node.infer(cur_node)
  70
  71         out_node_size = len(cur_node.out_nodes())
  72         for ind in range(out_node_size):
  73             node = cur_node.out_node(ind)
  74             q.append(node)
  75
  76
  77 def add_convolution_to_swap_xy_coordinates(graph: Graph, input_node: Node, coordinates_size: int):
  78     """
  79     The function add convolution node after the node 'input_node' to swap xy coordinates of the boxes produced
  80     by the node 'input_node'. It is expected that box coordinates are located in the fastest changing dimension of the
  81     'input_node' output, i.e. the input tensor could be reshaped to [num_boxes, 4] or [num_boxes, 5]. If the size is 5,
  82     then the 0-th element for each of num_boxes blocks is not changed and element 1 is swapped with element 2, element 3
  83     is swapped with element 4. This is the case when boxes coordinates are produced by the layer "Proposal". The exact
  84     amount of elements in each block is equal to the 'coordinates_size' parameter.
  85     :param graph: graph to operate on.
  86     :param input_node: node producing boxes coordinates.
  87     :param coordinates_size: integer value equal to 4 or 5.
  88     :return convolution node that swaps coordinates.
  89     """
  90     # swap of input tensor with 4 or 5 numbers describing boxes are supported
  91     assert (coordinates_size in [4, 5])
  92     input_reshape_4d_op = Reshape(input_node.graph, dict(dim=np.array([-1, 1, 1, coordinates_size])))
  93     input_reshape_4d_node = input_reshape_4d_op.create_node([input_node], dict(name=input_node.name + '/reshape_4d'))
  94     update_attrs(input_reshape_4d_node, 'shape_attrs', 'dim')
  95
  96     if coordinates_size == 5:
  97         # zero indexed element is not box coordinate ("batch id" in case of Proposal)
  98         conv_filter_data = np.array(np.array([[[[1, 0, 0, 0, 0],
  99                                                 [0, 0, 1, 0, 0],
 100                                                 [0, 1, 0, 0, 0],
 101                                                 [0, 0, 0, 0, 1],
 102                                                 [0, 0, 0, 1, 0]]]],
 103                                              dtype=np.float32))
 104     else:
 105         conv_filter_data = np.array(np.array([[[[0, 1, 0, 0],
 106                                                 [1, 0, 0, 0],
 107                                                 [0, 0, 0, 1],
 108                                                 [0, 0, 1, 0]]]],
 109                                              dtype=np.float32))
 110
 111     conv_filter_const_op = Const(graph, dict(value=conv_filter_data))
 112     conv_filter_const_node = conv_filter_const_op.create_node([], dict(name=input_node.name + '/weights'))
 113
 114     conv_op = Convolution(graph, {
 115         'bias_addable': True,
 116         'channel_dims': np.array([3]),
 117         'batch_dims': np.array([0]),
 118         'input_feature_channel': 2,
 119         'output_feature_channel': 3,
 120         'group': 1,
 121         'layout': 'NHWC',
 122     })
 123     return conv_op.create_node([input_reshape_4d_node, conv_filter_const_node], dict(name=input_node.name + "/conv"))
 124
 125
 126 def add_fake_background_loc(graph: Graph, input_node: Node):
 127     """
 128     DetectionOutput layer expects that box coordinates contains coordinates of boxes for the "background" class also,
 129     but in the TensorFlow\* Object Detection API the tensor contains information about real object classes only.
 130     The function copies a slice of the output data of the node 'input_node' and then concats it to the beginning of the
 131     data. The data in this slice is not used by the Detection Output layer so the actual values are not important. This
 132     approach allows the model to be reshape-able and does not introduce many layers.
 133     "background" class box coordinates.
 134     :param graph: graph to operate on.
 135     :param input_node: node producing the boxes coordinates.
 136     :return convolution node that adds slice of data for the "background" class.
 137     """
 138     crop_op = Crop(graph, dict(axis=np.array([1]), offset=np.array([0]), dim=np.array([1]), nchw_layout=True))
 139     crop_node = crop_op.create_node([input_node], dict(name='crop_locs'))
 140
 141     concat_op = Concat(graph, dict(axis=1, in_ports_count=2, nchw_layout=True))
 142     return concat_op.create_node([crop_node, input_node], dict(name=input_node.id + '/locs_with_fake_background'))
 143
 144
 145 def add_activation_function_after_node(graph: Graph, node: Node, activation_function: str):
 146     """
 147     The function adds node with activation function defined by string 'activation_function' which gets input from the
 148     node 'node'.
 149     :param graph: graph to operate on.
 150     :param node: node to add activation after.
 151     :param activation_function: string defining the activation function. These values are read from TensorFlow* object
 152     detection API pipeline configuration file
 153     :return: activation function node.
 154     """
 155     if activation_function == 'SOFTMAX':
 156         # softmax to be applied to the confidence
 157         softmax_conf_op = Softmax(graph, dict(axis=-1, nchw_layout=True))
 158         activation_node = softmax_conf_op.create_node([node], dict(name=node.name + '/softmax'))
 159     elif activation_function == 'SIGMOID':
 160         # sigmoid activation function to be applied to the confidence
 161         sigmoid_conf_op = Activation(graph, dict(operation='sigmoid', nchw_layout=True))
 162         activation_node = sigmoid_conf_op.create_node([node], dict(name=node.name + '/sigmoid'))
 163     elif activation_function == 'IDENTITY':
 164         # in case of Identity do nothing and just use result from the input node
 165         activation_node = node
 166     else:
 167         raise Error('Unknown post-processing activation function "{}".'.format(activation_function))
 168     return activation_node