diff options
Diffstat (limited to 'src/armnn')
-rw-r--r-- | src/armnn/layers/BatchToSpaceNdLayer.cpp | 35 | ||||
-rw-r--r-- | src/armnn/test/LayerValidateOutputTest.cpp | 38 |
2 files changed, 71 insertions, 2 deletions
diff --git a/src/armnn/layers/BatchToSpaceNdLayer.cpp b/src/armnn/layers/BatchToSpaceNdLayer.cpp index 9366a8710b..aff818e664 100644 --- a/src/armnn/layers/BatchToSpaceNdLayer.cpp +++ b/src/armnn/layers/BatchToSpaceNdLayer.cpp @@ -66,10 +66,41 @@ std::vector<TensorShape> BatchToSpaceNdLayer::InferOutputShapes(const std::vecto std::pair<unsigned int, unsigned int> xCrops = crops[1]; unsigned int inputHeight = inputShape[dataLayout.GetHeightIndex()]; - unsigned int outputHeight = theBlockShape.at(0) * (inputHeight - (yCrops.first + yCrops.second)); + unsigned int outputHeight; + unsigned int yCropsTotal = yCrops.first + yCrops.second; + + BOOST_ASSERT_MSG(yCropsTotal <= inputHeight, + "BatchToSpaceLayer: Overall height crop should be less than or equal to the input height."); + + unsigned int croppedHeight = inputHeight - yCropsTotal; + + if (theBlockShape.at(0) > 0) + { + outputHeight = theBlockShape.at(0) * croppedHeight; + } + else + { + outputHeight = croppedHeight; + } + + unsigned int outputWidth; unsigned int inputWidth = inputShape[dataLayout.GetWidthIndex()]; - unsigned int outputWidth = theBlockShape.at(1) * (inputWidth - (xCrops.first + xCrops.second)); + + unsigned int xCropsTotal = xCrops.first + xCrops.second; + + BOOST_ASSERT_MSG(xCropsTotal <= inputWidth, + "BatchToSpaceLayer: Overall width crop should be less than or equal to the input width."); + unsigned int croppedWidth = inputWidth - xCropsTotal; + + if (theBlockShape.at(1) > 0) + { + outputWidth = theBlockShape.at(1) * croppedWidth; + } + else + { + outputWidth = croppedWidth; + } unsigned int outputBatchSize = overallSize / (outputHeight * outputWidth); diff --git a/src/armnn/test/LayerValidateOutputTest.cpp b/src/armnn/test/LayerValidateOutputTest.cpp new file mode 100644 index 0000000000..62b9c4a0d8 --- /dev/null +++ b/src/armnn/test/LayerValidateOutputTest.cpp @@ -0,0 +1,38 @@ +// +// Copyright © 2017 Arm Ltd. All rights reserved. +// SPDX-License-Identifier: MIT +// +#include <armnn/ArmNN.hpp> + +#include <boost/algorithm/string.hpp> +#include <boost/test/unit_test.hpp> +#include <layers/BatchToSpaceNdLayer.hpp> +#include <Graph.hpp> + + +BOOST_AUTO_TEST_SUITE(LayerValidateOutput) + +BOOST_AUTO_TEST_CASE(TestBatchToSpaceInferOutputShape) +{ + armnn::Graph graph; + + armnn::BatchToSpaceNdDescriptor descriptor; + std::vector<unsigned int> theBlockShape = {2, 2}; + descriptor.m_BlockShape = theBlockShape; + descriptor.m_DataLayout = armnn::DataLayout::NHWC; + + armnn::BatchToSpaceNdLayer* const batchToSpaceLayer = + graph.AddLayer<armnn::BatchToSpaceNdLayer>(descriptor, "batchToSpace"); + + std::vector<armnn::TensorShape> shapes; + const std::vector<unsigned int> theDimSizes = {4, 2, 2, 1}; + armnn::TensorShape shape(4, theDimSizes.data()); + shapes.push_back(shape); + + const std::vector<unsigned int> expectedDimSizes = {1, 4, 4, 1}; + armnn::TensorShape expectedShape(4, expectedDimSizes.data()); + + BOOST_CHECK(expectedShape == batchToSpaceLayer->InferOutputShapes(shapes).at(0)); +} + +BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file |