model-optimizer/extensions/front/onnx/upsample_ext.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 logging as log
  18
  19 from extensions.ops.resample import ResampleOp
  20 from mo.front.extractor import FrontExtractorOp
  21 from mo.front.onnx.extractors.utils import onnx_attr
  22 from mo.utils.error import Error
  23
  24 import numpy as np
  25
  26
  27 class UpsampleFrontExtractor(FrontExtractorOp):
  28     op = 'Upsample'
  29     enabled = True
  30
  31     @staticmethod
  32     def extract(node):
  33         mode = onnx_attr(node, 'mode', 's', default='nearest', dst_type=lambda x: x.decode())
  34         scales = onnx_attr(node, 'scales', 'floats', dst_type=lambda x: np.array(x, dtype=np.float32))
  35         width_scale = onnx_attr(node, 'width_scale', 'f')
  36         height_scale = onnx_attr(node, 'height_scale', 'f')
  37
  38         supported_modes = ['nearest', 'linear']
  39         if mode not in supported_modes:
  40             raise Error(
  41                 'Error decoding Upsample node {}, mode = {} is not in the list of supported modes {}.',
  42                 node.name,
  43                 mode,
  44                 supported_modes
  45             )
  46
  47         # TODO: this is a temporary limitation
  48         if mode != 'nearest':
  49             raise Error(
  50                 'Upsample mode {} for node {} is not supported. Only nearest is supported.',
  51                 mode,
  52                 node.name
  53             )
  54
  55         # TODO: this is a temporary limitation
  56         if scales is not None:
  57             raise Error(
  58                 'Upsample scales attribute is defined for node {}. Only scale_width and scale_height are supported.',
  59                 node.name
  60             )
  61
  62         if width_scale is None or height_scale is None:
  63             raise Error(
  64                 'One/both of widths_scale = {} and height_scale = {} is not defined for Upsampe node {}.',
  65                 width_scale,
  66                 height_scale,
  67                 node.name
  68             )
  69
  70         if width_scale != height_scale:
  71             raise Error(
  72                 'Upsample node {} have different widths_scale = {} and height_scale = {}. It is not supported; they should match.',
  73                 node.name,
  74                 width_scale,
  75                 height_scale
  76             )
  77
  78         mode_to_resample_type = {'nearest': 'caffe.ResampleParameter.NEAREST'}
  79         assert mode in mode_to_resample_type
  80         assert width_scale == height_scale
  81         assert width_scale is not None
  82         ResampleOp.update_node_stat(node, {'resample_type': mode_to_resample_type[mode], 'factor': width_scale, 'antialias': 0})
  83         return __class__.enabled