1 //
2 // Copyright © 2017 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 #pragma once
6
7 #include <ResolveType.hpp>
8
9 #include <armnn/INetwork.hpp>
10
11 #include <CommonTestUtils.hpp>
12
13 #include <doctest/doctest.h>
14
15 #include <vector>
16
17 namespace
18 {
19
20 template<typename armnn::DataType DataType>
CreateBatchToSpaceNdNetwork(const armnn::TensorShape & inputShape,const armnn::TensorShape & outputShape,std::vector<unsigned int> & blockShape,std::vector<std::pair<unsigned int,unsigned int>> & crops,armnn::DataLayout dataLayout,const float qScale=1.0f,const int32_t qOffset=0)21 INetworkPtr CreateBatchToSpaceNdNetwork(const armnn::TensorShape& inputShape,
22 const armnn::TensorShape& outputShape,
23 std::vector<unsigned int>& blockShape,
24 std::vector<std::pair<unsigned int, unsigned int>>& crops,
25 armnn::DataLayout dataLayout,
26 const float qScale = 1.0f,
27 const int32_t qOffset = 0)
28 {
29 using namespace armnn;
30 // Builds up the structure of the network.
31 INetworkPtr net(INetwork::Create());
32
33 TensorInfo inputTensorInfo(inputShape, DataType, qScale, qOffset, true);
34 TensorInfo outputTensorInfo(outputShape, DataType, qScale, qOffset);
35
36 BatchToSpaceNdDescriptor batchToSpaceNdDesc(blockShape, crops);
37 batchToSpaceNdDesc.m_DataLayout = dataLayout;
38
39 IConnectableLayer* batchToSpaceNd = net->AddBatchToSpaceNdLayer(batchToSpaceNdDesc, "batchToSpaceNd");
40 IConnectableLayer* input = net->AddInputLayer(0, "input");
41 IConnectableLayer* output = net->AddOutputLayer(0, "output");
42
43 Connect(batchToSpaceNd, output, outputTensorInfo, 0, 0);
44 Connect(input, batchToSpaceNd, inputTensorInfo, 0, 0);
45
46 return net;
47 }
48
49 template<armnn::DataType ArmnnType>
BatchToSpaceNdEndToEnd(const std::vector<BackendId> & backends,armnn::DataLayout dataLayout)50 void BatchToSpaceNdEndToEnd(const std::vector<BackendId>& backends, armnn::DataLayout dataLayout)
51 {
52 using namespace armnn;
53 using T = ResolveType<ArmnnType>;
54
55 std::vector<unsigned int> blockShape {2, 2};
56 std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
57 const TensorShape& inputShape = { 4, 1, 1, 1 };
58 const TensorShape& outputShape = (dataLayout == DataLayout::NCHW)
59 ? std::initializer_list<unsigned int>({ 1, 1, 2, 2 })
60 : std::initializer_list<unsigned int>({ 1, 2, 2, 1 });
61
62 // Builds up the structure of the network
63 INetworkPtr net = CreateBatchToSpaceNdNetwork<ArmnnType>(inputShape, outputShape, blockShape, crops, dataLayout);
64
65 CHECK(net);
66
67 // Creates structures for input & output.
68 std::vector<T> inputData{ 1, 2, 3, 4 };
69
70 std::vector<T> expectedOutput{ 1, 2, 3, 4 };
71
72 std::map<int, std::vector<T>> inputTensorData = { { 0, inputData } };
73 std::map<int, std::vector<T>> expectedOutputData = { { 0, expectedOutput } };
74
75 EndToEndLayerTestImpl<ArmnnType, ArmnnType>(move(net), inputTensorData, expectedOutputData, backends);
76 }
77
78 template<armnn::DataType ArmnnType>
BatchToSpaceNdComplexEndToEnd(const std::vector<BackendId> & backends,armnn::DataLayout dataLayout)79 void BatchToSpaceNdComplexEndToEnd(const std::vector<BackendId>& backends, armnn::DataLayout dataLayout)
80 {
81 using namespace armnn;
82 using T = ResolveType<ArmnnType>;
83
84 std::vector<unsigned int> blockShape {2, 2};
85 std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {2, 0}};
86 const TensorShape& inputShape = (dataLayout == DataLayout::NCHW)
87 ? std::initializer_list<unsigned int>({ 8, 1, 1, 3 })
88 : std::initializer_list<unsigned int>({ 8, 1, 3, 1 });
89 const TensorShape& outputShape = (dataLayout == DataLayout::NCHW)
90 ? std::initializer_list<unsigned int>({ 2, 1, 2, 4 })
91 : std::initializer_list<unsigned int>({ 2, 2, 4, 1 });
92
93 // Builds up the structure of the network
94 INetworkPtr net = CreateBatchToSpaceNdNetwork<ArmnnType>(inputShape, outputShape, blockShape, crops, dataLayout);
95
96 CHECK(net);
97
98 // Creates structures for input & output.
99 std::vector<T> inputData{
100 0, 1, 3, 0, 9, 11,
101 0, 2, 4, 0, 10, 12,
102 0, 5, 7, 0, 13, 15,
103 0, 6, 8, 0, 14, 16
104 };
105
106 std::vector<T> expectedOutput{
107 1, 2, 3, 4,
108 5, 6, 7, 8,
109 9, 10, 11, 12,
110 13, 14, 15, 16
111 };
112
113 std::map<int, std::vector<T>> inputTensorData = { { 0, inputData } };
114 std::map<int, std::vector<T>> expectedOutputData = { { 0, expectedOutput } };
115
116 EndToEndLayerTestImpl<ArmnnType, ArmnnType>(move(net), inputTensorData, expectedOutputData, backends);
117 }
118
119 } // anonymous namespace
120