model-optimizer/mo/ops/slice.py

   1 """
   2  Copyright (c) 2018 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 copy
  18
  19 import networkx as nx
  20 import numpy as np
  21 from mo.graph.graph import Node
  22 from mo.ops.op import Op
  23 import logging as log
  24
  25
  26 class Slice(Op):
  27     op = 'Slice'
  28     enabled = True
  29
  30     def __init__(self, graph: nx.MultiDiGraph, attrs: dict):
  31         super().__init__(graph, {
  32             'op': 'Slice',
  33             'infer': __class__.infer
  34         }, attrs)
  35
  36     @staticmethod
  37     def infer(node: Node):
  38         if node.start is None or node.end is None:
  39             log.warning('Incorrect slice operation: no starts or ends attr')
  40             return
  41
  42         if len(node.in_nodes()) != 1:
  43             log.warning('Incorrect slice operation: slice op should have exactly one input')
  44             return
  45
  46         if node.in_node(0).value is None:
  47             log.info('Slice operation supports only on constant path')
  48             return
  49
  50         axis = node.axis
  51         start = node.start
  52         end = node.end
  53         value = node.in_node(0).value
  54         input_shape = node.in_node(0).shape
  55
  56         # Following ONNX specification, in case of unknown axis, axises should be in greater order
  57         if axis is None:
  58             axis = [x for x in range(len(start))]
  59
  60         # Calculate output value for slice operation
  61         slice_idx = [None for x in range(len(axis))]
  62         for id in range(len(axis)):
  63             # Ranged for output value for specified axis
  64             slice_idx[axis[id]] = slice(start[id], end[id], 1)
  65
  66         for axis, s in enumerate(slice_idx):
  67             if s is None:
  68                 slice_idx[axis] = slice(0, input_shape[axis], 1)
  69
  70         value = value[slice_idx]
  71         node.out_node().value = np.array(value)
  72         node.out_node().shape = np.array(value.shape)