diff options
Diffstat (limited to 'src/armnn/test/InferOutputTests.hpp')
-rw-r--r-- | src/armnn/test/InferOutputTests.hpp | 86 |
1 files changed, 86 insertions, 0 deletions
diff --git a/src/armnn/test/InferOutputTests.hpp b/src/armnn/test/InferOutputTests.hpp index c428a9db61..feb2125656 100644 --- a/src/armnn/test/InferOutputTests.hpp +++ b/src/armnn/test/InferOutputTests.hpp @@ -10,6 +10,7 @@ #include <armnn/ArmNN.hpp> #include <Graph.hpp> +#include <layers/ArgMinMaxLayer.hpp> #include <layers/BatchToSpaceNdLayer.hpp> #include <layers/SpaceToDepthLayer.hpp> #include <layers/PreluLayer.hpp> @@ -18,6 +19,91 @@ #include <boost/algorithm/string.hpp> #include <boost/test/unit_test.hpp> +void ArgMinMaxInferOutputShapeImpl(const armnn::ArgMinMaxDescriptor descriptor, + const std::vector<armnn::TensorShape>& inputShapes, + std::vector<armnn::TensorShape>& outputShapes) +{ + armnn::Graph graph; + auto argMinMaxLayer = graph.AddLayer<armnn::ArgMinMaxLayer>(descriptor, "argMinMax"); + outputShapes = argMinMaxLayer->InferOutputShapes(inputShapes); +} + +void ArgMinMaxInferOutputShape4dTest() +{ + armnn::Graph graph; + armnn::ArgMinMaxDescriptor descriptor; + descriptor.m_Axis = 2; + + const std::vector<armnn::TensorShape> inputShapes + { + { 1, 3, 2, 4 } + }; + + std::vector<armnn::TensorShape> outputShapes; + BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes)); + + armnn::TensorShape expectedOutputShape( { 1, 3, 4 } ); + BOOST_CHECK(outputShapes.size() == 1); + BOOST_CHECK(outputShapes[0] == expectedOutputShape); +} + +void ArgMinMaxInferOutputShape3dTest() +{ + armnn::Graph graph; + armnn::ArgMinMaxDescriptor descriptor; + descriptor.m_Axis = 0; + + const std::vector<armnn::TensorShape> inputShapes + { + { 1, 3, 2 } + }; + + std::vector<armnn::TensorShape> outputShapes; + BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes)); + + armnn::TensorShape expectedOutputShape( { 3, 2 } ); + BOOST_CHECK(outputShapes.size() == 1); + BOOST_CHECK(outputShapes[0] == expectedOutputShape); +} + +void ArgMinMaxInferOutputShape2dTest() +{ + armnn::Graph graph; + armnn::ArgMinMaxDescriptor descriptor; + descriptor.m_Axis = 1; + + const std::vector<armnn::TensorShape> inputShapes + { + { 3, 2 } + }; + + std::vector<armnn::TensorShape> outputShapes; + BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes)); + + armnn::TensorShape expectedOutputShape( { 3 } ); + BOOST_CHECK(outputShapes.size() == 1); + BOOST_CHECK(outputShapes[0] == expectedOutputShape); +} + +void ArgMinMaxInferOutputShape1dTest() +{ + armnn::Graph graph; + armnn::ArgMinMaxDescriptor descriptor; + descriptor.m_Axis = 0; + + const std::vector<armnn::TensorShape> inputShapes + { + { 5 } + }; + + std::vector<armnn::TensorShape> outputShapes; + BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes)); + + armnn::TensorShape expectedOutputShape( { 1 } ); + BOOST_CHECK(outputShapes.size() == 1); + BOOST_CHECK(outputShapes[0] == expectedOutputShape); +} + void BatchToSpaceInferOutputShapeTest() { armnn::Graph graph; |