model-optimizer/mo/front/caffe/loader_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 unittest
  18
  19 import numpy as np
  20 from google.protobuf import text_format
  21
  22 from mo.front.caffe.loader import caffe_pb_to_nx
  23 from mo.front.caffe.proto import caffe_pb2
  24 from mo.utils.error import Error
  25
  26 proto_str_one_input = 'name: "network" ' \
  27                       'layer { ' \
  28                       'name: "Input0" ' \
  29                       'type: "Input" ' \
  30                       'top: "Input0" ' \
  31                       'input_param { ' \
  32                       'shape: { ' \
  33                       'dim: 1 ' \
  34                       'dim: 3 ' \
  35                       'dim: 224 ' \
  36                       'dim: 224 ' \
  37                       '} ' \
  38                       '} ' \
  39                       '}'
  40
  41 proto_str_old_styled_multi_input = 'name: "network" ' \
  42                                    'input: "Input0" ' \
  43                                    'input_dim: 1 ' \
  44                                    'input_dim: 3 ' \
  45                                    'input_dim: 224 ' \
  46                                    'input_dim: 224 ' \
  47                                    'input: "data"' \
  48                                    'input_dim: 1 ' \
  49                                    'input_dim: 3 '
  50
  51 proto_str_input = 'name: "network" ' \
  52                   'input: "data" ' \
  53                   'input_shape ' \
  54                   '{ ' \
  55                   'dim: 1 ' \
  56                   'dim: 3 ' \
  57                   'dim: 224 ' \
  58                   'dim: 224 ' \
  59                   '}'
  60
  61 proto_str_multi_input = 'name: "network" ' \
  62                         'input: "data" ' \
  63                         'input_shape ' \
  64                         '{ ' \
  65                         'dim: 1 ' \
  66                         'dim: 3 ' \
  67                         'dim: 224 ' \
  68                         'dim: 224 ' \
  69                         '} ' \
  70                         'input: "data1"' \
  71                         'input_shape ' \
  72                         '{ ' \
  73                         'dim: 1 ' \
  74                         'dim: 3 ' \
  75                         '}'
  76
  77 proto_str_old_styled_input = 'name: "network" ' \
  78                              'input: "data" ' \
  79                              'input_dim: 1 ' \
  80                              'input_dim: 3 ' \
  81                              'input_dim: 224 ' \
  82                              'input_dim: 224 '
  83
  84 layer_proto_str = 'layer { ' \
  85                   'name: "conv1" ' \
  86                   'type: "Convolution" ' \
  87                   'bottom: "data" ' \
  88                   'top: "conv1" ' \
  89                   '}'
  90
  91 proto_same_name_layers = 'layer { ' \
  92                          'name: "conv1" ' \
  93                          'type: "Convolution" ' \
  94                          'bottom: "data" ' \
  95                          'top: "conv1" ' \
  96                          '}' \
  97                          'layer { ' \
  98                          'name: "conv1" ' \
  99                          'type: "Convolution" ' \
 100                          'bottom: "data1" ' \
 101                          'top: "conv1_2" ' \
 102                          '}'
 103
 104 class TestLoader(unittest.TestCase):
 105     def test_caffe_pb_to_nx_one_input(self):
 106         proto = caffe_pb2.NetParameter()
 107         text_format.Merge(proto_str_one_input, proto)
 108         graph, input_shapes = caffe_pb_to_nx(proto, None)
 109         expected_input_shapes = {
 110             'Input0': np.array([1, 3, 224, 224])
 111         }
 112
 113         for i in expected_input_shapes:
 114             np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
 115
 116     def test_caffe_pb_to_nx_old_styled_multi_input(self):
 117         proto = caffe_pb2.NetParameter()
 118         text_format.Merge(proto_str_old_styled_multi_input + layer_proto_str, proto)
 119         self.assertRaises(Error, caffe_pb_to_nx, proto, None)
 120
 121     def test_caffe_pb_to_nx_old_styled_input(self):
 122         proto = caffe_pb2.NetParameter()
 123         text_format.Merge(proto_str_old_styled_input + layer_proto_str, proto)
 124         graph, input_shapes = caffe_pb_to_nx(proto, None)
 125         expected_input_shapes = {
 126             'data': np.array([1, 3, 224, 224])
 127         }
 128
 129         for i in expected_input_shapes:
 130             np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
 131
 132     def test_caffe_pb_to_standart_input(self):
 133         proto = caffe_pb2.NetParameter()
 134         text_format.Merge(proto_str_input + layer_proto_str, proto)
 135         graph, input_shapes = caffe_pb_to_nx(proto, None)
 136         expected_input_shapes = {
 137             'data': np.array([1, 3, 224, 224])
 138         }
 139
 140         for i in expected_input_shapes:
 141             np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
 142
 143     def test_caffe_pb_to_multi_input(self):
 144         proto = caffe_pb2.NetParameter()
 145         text_format.Merge(proto_str_multi_input + layer_proto_str, proto)
 146         graph, input_shapes = caffe_pb_to_nx(proto, None)
 147         expected_input_shapes = {
 148             'data': np.array([1, 3, 224, 224]),
 149             'data1': np.array([1, 3])
 150         }
 151
 152         for i in expected_input_shapes:
 153             np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
 154
 155     def test_caffe_same_name_layer(self):
 156         proto = caffe_pb2.NetParameter()
 157         text_format.Merge(proto_str_multi_input + proto_same_name_layers, proto)
 158         graph, input_shapes = caffe_pb_to_nx(proto, None)
 159         # 6 nodes because: 2 inputs + 2 convolutions
 160         np.testing.assert_equal(len(graph.nodes()), 4)