diff options
author | narpra01 <narumol.prangnawarat@arm.com> | 2018-10-26 17:36:32 +0100 |
---|---|---|
committer | Matthew Bentham <matthew.bentham@arm.com> | 2018-10-27 21:50:38 +0000 |
commit | 7af7688a77e9fe2929a626f7c7dab9aef089d2b3 (patch) | |
tree | 70ab471953ee151107b411588a515a25ca845b97 /src/armnn/layers/Pooling2dLayer.cpp | |
parent | f98d21a244c384c5942e4b261d168f65ecc8a797 (diff) | |
download | armnn-7af7688a77e9fe2929a626f7c7dab9aef089d2b3.tar.gz |
IVGCVSW-2083 - Add DataLayout parameter to calculate the expected output shape in ValidateTensorShapesFromInputs
* Convolution2dLayer
* DepthwiseConvolution2dLayer
* Pooling2dLayer
* ResizeBilinearLayer
* Unittests for ValidateTensorShapesFromInputs
Change-Id: I057caf8a90d822175a7dd7271f960b65c6154bb4
Diffstat (limited to 'src/armnn/layers/Pooling2dLayer.cpp')
-rw-r--r-- | src/armnn/layers/Pooling2dLayer.cpp | 13 |
1 files changed, 8 insertions, 5 deletions
diff --git a/src/armnn/layers/Pooling2dLayer.cpp b/src/armnn/layers/Pooling2dLayer.cpp index d87ad0f19f..779ac2041e 100644 --- a/src/armnn/layers/Pooling2dLayer.cpp +++ b/src/armnn/layers/Pooling2dLayer.cpp @@ -37,10 +37,9 @@ std::vector<TensorShape> Pooling2dLayer::InferOutputShapes(const std::vector<Ten // If we support multiple batch dimensions in the future, then this assert will need to change. BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Pooling2dLayer will always have 4D input."); - - unsigned int inWidth = inputShape[3]; - unsigned int inHeight = inputShape[2]; - unsigned int inChannels = inputShape[1]; + unsigned int inWidth = inputShape[m_Param.m_DataLayout.GetWidthIndex()]; + unsigned int inHeight = inputShape[m_Param.m_DataLayout.GetHeightIndex()]; + unsigned int inChannels = inputShape[m_Param.m_DataLayout.GetChannelsIndex()]; unsigned int inBatchSize = inputShape[0]; bool isGlobalPooling = (m_Param.m_StrideX==0 && m_Param.m_StrideY==0); @@ -88,7 +87,11 @@ std::vector<TensorShape> Pooling2dLayer::InferOutputShapes(const std::vector<Ten unsigned int outChannels = inChannels; unsigned int outBatchSize = inBatchSize; - return std::vector<TensorShape>({ TensorShape({outBatchSize, outChannels, outHeight, outWidth}) }); + TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ? + TensorShape( { outBatchSize, outHeight, outWidth, outChannels } ) : + TensorShape( { outBatchSize, outChannels, outHeight, outWidth }); + + return std::vector<TensorShape>({ tensorShape }); } void Pooling2dLayer::ValidateTensorShapesFromInputs() |