diff options
Diffstat (limited to 'src/armnn/test/LayerValidateOutputTest.cpp')
| -rw-r--r-- | src/armnn/test/LayerValidateOutputTest.cpp | 29 |
1 files changed, 26 insertions, 3 deletions
diff --git a/src/armnn/test/LayerValidateOutputTest.cpp b/src/armnn/test/LayerValidateOutputTest.cpp index 999844e252..acefd51110 100644 --- a/src/armnn/test/LayerValidateOutputTest.cpp +++ b/src/armnn/test/LayerValidateOutputTest.cpp @@ -4,11 +4,12 @@ // #include <armnn/ArmNN.hpp> -#include <boost/algorithm/string.hpp> -#include <boost/test/unit_test.hpp> -#include <layers/BatchToSpaceNdLayer.hpp> #include <Graph.hpp> +#include <layers/BatchToSpaceNdLayer.hpp> +#include <layers/SpaceToDepthLayer.hpp> +#include <boost/algorithm/string.hpp> +#include <boost/test/unit_test.hpp> BOOST_AUTO_TEST_SUITE(LayerValidateOutput) @@ -35,4 +36,26 @@ BOOST_AUTO_TEST_CASE(TestBatchToSpaceInferOutputShape) BOOST_CHECK(expectedShape == batchToSpaceLayer->InferOutputShapes(shapes).at(0)); } +BOOST_AUTO_TEST_CASE(TestSpaceToDepthInferOutputShape) +{ + armnn::Graph graph; + + armnn::SpaceToDepthDescriptor descriptor; + descriptor.m_BlockSize = 2; + descriptor.m_DataLayout = armnn::DataLayout::NHWC; + + armnn::SpaceToDepthLayer* const spaceToDepthLayer = + graph.AddLayer<armnn::SpaceToDepthLayer>(descriptor, "spaceToDepth"); + + std::vector<armnn::TensorShape> shapes; + const std::vector<unsigned int> dimSizes{ 1, 16, 8, 3 }; + armnn::TensorShape shape(4, dimSizes.data()); + shapes.push_back(shape); + + const std::vector<unsigned int> expectedDimSizes{ 1, 8, 4, 12 }; + armnn::TensorShape expectedShape(4, expectedDimSizes.data()); + + BOOST_CHECK(expectedShape == spaceToDepthLayer->InferOutputShapes(shapes).at(0)); +} + BOOST_AUTO_TEST_SUITE_END() |