model-optimizer/mo/front/onnx/extractors/constant_test.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 import unittest
  19
  20 import numpy as np
  21 import onnx
  22 from generator import generator, generate
  23 from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE
  24
  25 from mo.front.onnx.extractors.constant import onnx_constant_ext
  26 from mo.utils.unittest.extractors import PB
  27
  28 dtypes = [np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.float32, np.double, np.bool]
  29
  30
  31 @generator
  32 class ConstantONNXExtractorTest(unittest.TestCase):
  33     @staticmethod
  34     def _create_constant_node(numpy_dtype):
  35         numpy_dtype = np.dtype(numpy_dtype)
  36         if numpy_dtype not in NP_TYPE_TO_TENSOR_TYPE:
  37             log.error("Numpy type {} not supported in ONNX".format(numpy_dtype))
  38             return None
  39
  40         values = np.array(np.random.randn(5, 5).astype(numpy_dtype))
  41         pb = onnx.helper.make_node(
  42             'Constant',
  43             inputs=[],
  44             outputs=['values'],
  45             value=onnx.helper.make_tensor(
  46                 name='const_tensor',
  47                 data_type=NP_TYPE_TO_TENSOR_TYPE[numpy_dtype],
  48                 dims=values.shape,
  49                 vals=values.flatten().astype(numpy_dtype),
  50             ),
  51         )
  52         node = PB({'pb': pb})
  53         return node
  54
  55     @generate(*dtypes)
  56     def test_constant_ext(self, np_dtype):
  57         node = self._create_constant_node(np_dtype)
  58         attrs = onnx_constant_ext(node)
  59         self.assertTrue(attrs['data_type'] == np_dtype,
  60                         'Wrong data_type attribute: recieved {}, expected {}'.format(attrs['data_type'], np_dtype))