diff options
author | Éanna Ó Catháin <eanna.ocathain@arm.com> | 2018-11-14 11:26:23 +0000 |
---|---|---|
committer | Matteo Martincigh <matteo.martincigh@arm.com> | 2018-11-14 12:00:22 +0000 |
commit | 262553e0d4e2571c8d7b3c57a7271c23a783e04b (patch) | |
tree | 05fa2b4b572a19e2de8eb3455f348416aca1eb0f /src/armnn/layers/BatchToSpaceNdLayer.cpp | |
parent | 3aab7c301bcbc206169c12ac04162b1445b4d472 (diff) | |
download | armnn-262553e0d4e2571c8d7b3c57a7271c23a783e04b.tar.gz |
IVGCVSW-2054: Fixing issue with InferOutputShape implementation in BatchToSpaceNdLayer.
* added Unit test
Change-Id: I80f55e8c7afb39d96006c8dd027fc9683ea8182e
Diffstat (limited to 'src/armnn/layers/BatchToSpaceNdLayer.cpp')
-rw-r--r-- | src/armnn/layers/BatchToSpaceNdLayer.cpp | 35 |
1 files changed, 33 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); |