161 lines
5.9 KiB
Python
161 lines
5.9 KiB
Python
"""
|
|
Copyright (c) 2018-2019 Intel Corporation
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from google.protobuf import text_format
|
|
|
|
from mo.front.caffe.loader import caffe_pb_to_nx
|
|
from mo.front.caffe.proto import caffe_pb2
|
|
from mo.utils.error import Error
|
|
|
|
proto_str_one_input = 'name: "network" ' \
|
|
'layer { ' \
|
|
'name: "Input0" ' \
|
|
'type: "Input" ' \
|
|
'top: "Input0" ' \
|
|
'input_param { ' \
|
|
'shape: { ' \
|
|
'dim: 1 ' \
|
|
'dim: 3 ' \
|
|
'dim: 224 ' \
|
|
'dim: 224 ' \
|
|
'} ' \
|
|
'} ' \
|
|
'}'
|
|
|
|
proto_str_old_styled_multi_input = 'name: "network" ' \
|
|
'input: "Input0" ' \
|
|
'input_dim: 1 ' \
|
|
'input_dim: 3 ' \
|
|
'input_dim: 224 ' \
|
|
'input_dim: 224 ' \
|
|
'input: "data"' \
|
|
'input_dim: 1 ' \
|
|
'input_dim: 3 '
|
|
|
|
proto_str_input = 'name: "network" ' \
|
|
'input: "data" ' \
|
|
'input_shape ' \
|
|
'{ ' \
|
|
'dim: 1 ' \
|
|
'dim: 3 ' \
|
|
'dim: 224 ' \
|
|
'dim: 224 ' \
|
|
'}'
|
|
|
|
proto_str_multi_input = 'name: "network" ' \
|
|
'input: "data" ' \
|
|
'input_shape ' \
|
|
'{ ' \
|
|
'dim: 1 ' \
|
|
'dim: 3 ' \
|
|
'dim: 224 ' \
|
|
'dim: 224 ' \
|
|
'} ' \
|
|
'input: "data1"' \
|
|
'input_shape ' \
|
|
'{ ' \
|
|
'dim: 1 ' \
|
|
'dim: 3 ' \
|
|
'}'
|
|
|
|
proto_str_old_styled_input = 'name: "network" ' \
|
|
'input: "data" ' \
|
|
'input_dim: 1 ' \
|
|
'input_dim: 3 ' \
|
|
'input_dim: 224 ' \
|
|
'input_dim: 224 '
|
|
|
|
layer_proto_str = 'layer { ' \
|
|
'name: "conv1" ' \
|
|
'type: "Convolution" ' \
|
|
'bottom: "data" ' \
|
|
'top: "conv1" ' \
|
|
'}'
|
|
|
|
proto_same_name_layers = 'layer { ' \
|
|
'name: "conv1" ' \
|
|
'type: "Convolution" ' \
|
|
'bottom: "data" ' \
|
|
'top: "conv1" ' \
|
|
'}' \
|
|
'layer { ' \
|
|
'name: "conv1" ' \
|
|
'type: "Convolution" ' \
|
|
'bottom: "data1" ' \
|
|
'top: "conv1_2" ' \
|
|
'}'
|
|
|
|
class TestLoader(unittest.TestCase):
|
|
def test_caffe_pb_to_nx_one_input(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_one_input, proto)
|
|
graph, input_shapes = caffe_pb_to_nx(proto, None)
|
|
expected_input_shapes = {
|
|
'Input0': np.array([1, 3, 224, 224])
|
|
}
|
|
|
|
for i in expected_input_shapes:
|
|
np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
|
|
|
|
def test_caffe_pb_to_nx_old_styled_multi_input(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_old_styled_multi_input + layer_proto_str, proto)
|
|
self.assertRaises(Error, caffe_pb_to_nx, proto, None)
|
|
|
|
def test_caffe_pb_to_nx_old_styled_input(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_old_styled_input + layer_proto_str, proto)
|
|
graph, input_shapes = caffe_pb_to_nx(proto, None)
|
|
expected_input_shapes = {
|
|
'data': np.array([1, 3, 224, 224])
|
|
}
|
|
|
|
for i in expected_input_shapes:
|
|
np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
|
|
|
|
def test_caffe_pb_to_standart_input(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_input + layer_proto_str, proto)
|
|
graph, input_shapes = caffe_pb_to_nx(proto, None)
|
|
expected_input_shapes = {
|
|
'data': np.array([1, 3, 224, 224])
|
|
}
|
|
|
|
for i in expected_input_shapes:
|
|
np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
|
|
|
|
def test_caffe_pb_to_multi_input(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_multi_input + layer_proto_str, proto)
|
|
graph, input_shapes = caffe_pb_to_nx(proto, None)
|
|
expected_input_shapes = {
|
|
'data': np.array([1, 3, 224, 224]),
|
|
'data1': np.array([1, 3])
|
|
}
|
|
|
|
for i in expected_input_shapes:
|
|
np.testing.assert_array_equal(input_shapes[i], expected_input_shapes[i])
|
|
|
|
def test_caffe_same_name_layer(self):
|
|
proto = caffe_pb2.NetParameter()
|
|
text_format.Merge(proto_str_multi_input + proto_same_name_layers, proto)
|
|
graph, input_shapes = caffe_pb_to_nx(proto, None)
|
|
# 6 nodes because: 2 inputs + 2 convolutions
|
|
np.testing.assert_equal(len(graph.nodes()), 4)
|