model-optimizer/mo/front/common/partial_infer/caffe_fallback_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 from unittest.mock import MagicMock
  19
  20 import numpy as np
  21
  22 from mo.front.common.partial_infer.caffe_fallback import build_net
  23 from mo.utils.unittest.extractors import FakeMultiParam, FakeValue
  24 from mo.utils.unittest.graph import build_graph
  25
  26
  27 class Net:
  28     def __init__(self, blobs):
  29         self.blobs = blobs
  30         self.reshape_blob = MagicMock(return_value=np.array([1, 1, 1, 1]))
  31         self.reshape = MagicMock(return_value=np.array([1, 1, 1, 1]))
  32         self.forward = MagicMock(return_value={'top_node': FakeValue(np.array([1, 3, 112, 112]))})
  33
  34
  35 my_mock_net = None
  36
  37
  38 class Caffe:
  39     def __init__(self):
  40         self.TEST = 'TEST'
  41
  42     def Net(self, *args):
  43         return my_mock_net
  44
  45
  46 class TestCaffeNativePartialInfer(unittest.TestCase):
  47     @classmethod
  48     def setUpClass(cls):
  49         import sys
  50         sys.modules['caffe'] = Caffe()
  51         cls.nodes_attributes = {
  52             'node_1': {'type': 'Input', 'kind': 'op'},
  53             'node_2': {'type': 'Input', 'kind': 'op'},
  54             'node_3': {'type': 'Identity', 'kind': 'op'},
  55             'node_4': {'type': 'Identity', 'kind': 'op'},
  56             'op_output': { 'kind': 'op', 'op': 'OpOutput'}
  57         }
  58
  59     def test_build_net_equal_inputs(self):
  60         global my_mock_net
  61         my_blobs = {
  62             'node_1': FakeValue(np.array([1, 3, 227, 227])),
  63             'node_2': FakeValue(np.array([1, 3, 224, 224]))
  64         }
  65         my_mock_net = Net(my_blobs)
  66         graph = build_graph(self.nodes_attributes,
  67                             [
  68                                 ('node_1', 'node_3'),
  69                                 ('node_2', 'node_3'),
  70                                 ('node_3', 'node_4'),
  71                                 ('node_4', 'op_output')
  72                             ],
  73                             {
  74                                 'node_4': {'shape': None},
  75                                 'node_1': {'shape': np.array([1, 3, 227, 227])},
  76                                 'node_2': {'shape': np.array([1, 3, 224, 224])},
  77                                 'node_3': {'top': 'top_node'}
  78                             })
  79         graph.proto_path = 'path_to_proto'
  80         graph.caffemodel_path = 'path_to_proto'
  81         build_net(graph)
  82         my_mock_net.reshape.assert_not_called()
  83         my_mock_net.forward.assert_called_once_with()
  84         self.assertIsNotNone(graph.caffe_net)
  85
  86     def test_build_net_not_equal_inputs(self):
  87         global my_mock_net
  88         input_node_param = {
  89             'shape': np.array([1, 3, 112, 112]),
  90             'reshape': MagicMock(return_value=134)
  91         }
  92         my_blobs = {
  93             'node_1': FakeMultiParam(input_node_param),
  94         }
  95         my_mock_net = Net(my_blobs)
  96         graph = build_graph(self.nodes_attributes,
  97                             [
  98                                 ('node_1', 'node_3'),
  99                                 ('node_3', 'node_4'),
 100                                 ('node_4', 'op_output')
 101                             ],
 102                             {'node_4': {'shape': None},
 103                              'node_1': {'shape': np.array([1, 3, 227, 227])},
 104                              'node_3': {'top': 'top_node'}
 105                              },
 106                             nodes_with_edges_only=True)
 107         graph.proto_path = 'path_to_proto'
 108         graph.caffemodel_path = 'path_to_proto'
 109         build_net(graph)
 110         my_mock_net.reshape.assert_called_once_with()
 111         my_mock_net.forward.assert_called_once_with()
 112         self.assertIsNotNone(graph.caffe_net)