def test_build_net_equal_inputs(self):
global my_mock_net
my_blobs = {
'node_1': FakeValue(np.array([1, 3, 227, 227])),
'node_2': FakeValue(np.array([1, 3, 224, 224]))
}
my_mock_net = Net(my_blobs)
graph = build_graph(
self.nodes_attributes, [('node_1', 'node_3'), ('node_2', 'node_3'),
('node_3', 'node_4'),
('node_4', 'op_output')], {
'node_4': {
'shape': None
},
'node_1': {
'shape': np.array([1, 3, 227, 227])
},
'node_2': {
'shape': np.array([1, 3, 224, 224])
},
'node_3': {
'top': 'top_node'
}
})
graph.proto_path = 'path_to_proto'
graph.caffemodel_path = 'path_to_proto'
build_net(graph)
my_mock_net.reshape.assert_not_called()
my_mock_net.forward.assert_called_once_with()
self.assertIsNotNone(graph.caffe_net)
def test_build_net_not_equal_inputs(self):
global my_mock_net
input_node_param = {
'shape': np.array([1, 3, 112, 112]),
'reshape': MagicMock(return_value=134)
}
my_blobs = {
'node_1': FakeMultiParam(input_node_param),
}
my_mock_net = Net(my_blobs)
graph = build_graph(self.nodes_attributes, [('node_1', 'node_3'),
('node_3', 'node_4'),
('node_4', 'op_output')],
{
'node_4': {
'shape': None
},
'node_1': {
'shape': np.array([1, 3, 227, 227])
},
'node_3': {
'top': 'top_node'
}
},
nodes_with_edges_only=True)
graph.proto_path = 'path_to_proto'
graph.caffemodel_path = 'path_to_proto'
build_net(graph)
my_mock_net.reshape.assert_called_once_with()
my_mock_net.forward.assert_called_once_with()
self.assertIsNotNone(graph.caffe_net)