model-optimizer/mo/ops/strided_slice.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
  17 import numpy as np
  18
  19 from mo.front.common.partial_infer.slice import tf_strided_slice_infer
  20 from mo.graph.graph import Node, Graph
  21 from mo.ops.op import Op, PermuteAttrs
  22 from mo.utils.utils import array_to_str
  23
  24
  25 def permute_array_with_ellipsis(node: Node, permutation: PermuteAttrs.Permutation, array: np.array, ins_value: int):
  26     """
  27     This function permutes masks according to permutation parameter. Several cases should be processed:
  28     * Some dimensions can be omitted in mask according to ellipsis mask
  29     * Mask length can be less than length of output dimensions plus shrinked dimensions
  30     * Mask have the same or more length than output
  31     """
  32     attr_mask_extended = list(array)
  33
  34     # If input and output have length of shape 3 and less, no need to permute
  35     if len(node.in_node().shape) < 4 and len(node.out_node().shape) < 4:
  36         return attr_mask_extended
  37
  38     # Length of mask is less than length of output ()plus shrinked dimensions then we should extend it before permutation
  39     if len(attr_mask_extended) < len(node.out_node(0).shape) + np.count_nonzero(node.shrink_axis_mask):
  40         # ellipsis is set, add dimensions in right place otherwise insert in the end
  41         if np.any(node.ellipsis_mask):
  42             idx = np.nonzero(node.ellipsis_mask)
  43             assert len(idx[0]) == 1
  44             id = idx[0][0]
  45         else:
  46             id = len(attr_mask_extended) - 1
  47
  48         ellips_ext = len(node.out_node(0).shape) + np.count_nonzero(node.shrink_axis_mask) - len(attr_mask_extended)
  49         for i in range(0, ellips_ext):
  50             attr_mask_extended.insert(id + i + 1, ins_value)
  51         # permute extended mask
  52         perm = PermuteAttrs.get_nhwc_to_nchw_permutation(len(attr_mask_extended))
  53         attr_mask_extended = np.array(attr_mask_extended)[perm.perm]
  54         return attr_mask_extended
  55     else:
  56         perm_len = len(node.out_node(0).shape) + np.count_nonzero(node.shrink_axis_mask)
  57         perm = PermuteAttrs.get_nhwc_to_nchw_permutation(perm_len)
  58         perm_list = list(perm.perm)
  59         # if mask length is more than output, just add tail that will not be permuted to avoid error
  60         for i in range(perm_len, len(attr_mask_extended)):
  61             perm_list.append(i)
  62         return np.array(attr_mask_extended, dtype=np.int64)[np.array(perm_list)]
  63
  64
  65 def permute_masks(node: Node, permutation: PermuteAttrs.Permutation, attr: str):
  66     if not node.has_valid(attr):
  67         return None
  68
  69     node[attr] = permute_array_with_ellipsis(node, permutation, node[attr],
  70                                              attr in ['begin_mask', 'end_mask'])
  71     return node[attr]
  72
  73
  74 class StridedSlice(Op):
  75     op = 'StridedSlice'
  76     enabled = True
  77
  78     def __init__(self, graph: Graph, attrs: dict):
  79         super().__init__(graph, {
  80             'type': __class__.op,
  81             'op': 'StridedSlice',
  82             'in_ports_count': 4,
  83             'out_ports_count': 1,
  84             'infer': __class__.infer
  85         }, attrs)
  86
  87     def backend_attrs(self):
  88         al = list()
  89
  90         def convert(attr):
  91             return lambda node: array_to_str(node, attr)
  92         for a in list(['new_axis_mask', 'shrink_axis_mask', 'ellipsis_mask', 'begin_mask', 'end_mask']):
  93             al.append((a, convert(a)))
  94         return al
  95
  96     @staticmethod
  97     def infer(node: Node):
  98         tf_strided_slice_infer(node)
  99
 100         PermuteAttrs.create_permute_attrs(node, attrs=[('shrink_axis_mask', 'input:0', permute_masks),
 101                                                        ('new_axis_mask', 'input:0', permute_masks),
 102                                                        ('ellipsis_mask', 'input:0', permute_masks),
 103                                                        ('begin_mask', 'input:0', permute_masks),
 104                                                        ('end_mask', 'input:0', permute_masks),
 105                                                        ])
 106
 107         for i in range(1, len(node.in_nodes())):
 108             if node.in_node(i).value is not None and node.in_node(i).shape[0] > 3:
 109                 perm = PermuteAttrs.get_nhwc_to_nchw_permutation(len(node.in_node(0).shape))
 110                 node.in_node(i).value = permute_array_with_ellipsis(node, perm, node.in_node(i).value, 0)
 111
 112         # due to permutation from nhwc to nchw we will extend all masks and inputs
 113         idx = np.nonzero(node.ellipsis_mask)
 114         node.ellipsis_mask[idx] = 0