model-optimizer/extensions/ops/accum.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 numpy as np
  18
  19 from mo.graph.graph import Node, Graph
  20 from mo.ops.op import Op
  21
  22
  23 class AccumOp(Op):
  24     op = 'Accum'
  25
  26     def __init__(self, graph: Graph, attrs: dict):
  27         mandatory_props = {
  28             'type': __class__.op,
  29             'op': __class__.op,
  30             'top_height': 0,
  31             'top_width': 0,
  32             'size_divisible_by': 0,
  33             'have_reference': 0,
  34             'out_ports_count': 1,
  35             'infer': AccumOp.accum_infer
  36         }
  37         super().__init__(graph, mandatory_props, attrs)
  38
  39     def supported_attrs(self):
  40         return [
  41             'top_height',
  42             'top_width',
  43             'size_divisible_by',
  44             'have_reference'
  45         ]
  46
  47     @staticmethod
  48     def accum_infer(node: Node):
  49
  50         batch = node.in_node(0).shape[0]
  51         num_inputs = len(node.in_nodes())
  52
  53         if node.have_reference:
  54             assert num_inputs >= 2, "Need at least two bottom blobs (one as reference)"
  55             total_channels = 0
  56             for i in range(num_inputs):
  57                 total_channels += node.in_node(i).shape[1]
  58                 assert node.in_node(i).shape[0] == batch, "All accumulated layers must have same number of images"
  59             assert total_channels >= 1, "Accumulated layers must have some channels in total"
  60             top_height_ = node.in_node(num_inputs - 1).shape[2]  # height
  61             top_width_ = node.in_node(num_inputs - 1).shape[3]  # width
  62             height_ = top_height_
  63             width_ = top_width_
  64         else:
  65             max_height = -1
  66             max_width = -1
  67             total_channels = 0
  68
  69             for i in range(num_inputs):
  70                 total_channels += node.in_node(i).shape[1]
  71                 max_height = node.in_node(i).shape[2] if node.in_node(i).shape[2] > max_height else max_height
  72                 max_width = node.in_node(i).shape[3] if node.in_node(i).shape[3] > max_width else max_width
  73                 assert node.in_node(i).shape[0] == batch, "All accumulated layers must have same number of images"
  74             assert total_channels >= 1, "Accumulated layers must have some channels in total"
  75
  76             if node.size_divisible_by:
  77                 sdb = node.size_divisible_by
  78                 top_height_ = int(np.ceil(max_height / sdb) * sdb)
  79                 top_width_ = int(np.ceil(max_width / sdb) * sdb)
  80             else:
  81                 top_height_ = node.top_height
  82                 top_width_ = node.top_width
  83             if top_height_ > max_height and top_width_ > max_width:  # Layer can specify custom top size which is larger than default
  84                 height_ = top_height_
  85                 width_ = top_width_
  86             else:  # Otherwise maximum of bottom sizes will be used
  87                 height_ = max_height
  88                 width_ = max_width
  89         channels_ = total_channels
  90         node.out_node(0).shape = np.array([batch, channels_, height_, width_])