aboutsummaryrefslogtreecommitdiff
path: root/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
diff options
context:
space:
mode:
authornarpra01 <narumol.prangnawarat@arm.com>2018-10-26 17:36:32 +0100
committerMatthew Bentham <matthew.bentham@arm.com>2018-10-27 21:50:38 +0000
commit7af7688a77e9fe2929a626f7c7dab9aef089d2b3 (patch)
tree70ab471953ee151107b411588a515a25ca845b97 /src/armnn/layers/DepthwiseConvolution2dLayer.cpp
parentf98d21a244c384c5942e4b261d168f65ecc8a797 (diff)
downloadarmnn-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/DepthwiseConvolution2dLayer.cpp')
-rw-r--r--src/armnn/layers/DepthwiseConvolution2dLayer.cpp22
1 files changed, 14 insertions, 8 deletions
diff --git a/src/armnn/layers/DepthwiseConvolution2dLayer.cpp b/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
index 393c4bf6f2..d80d3f1332 100644
--- a/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
+++ b/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
@@ -59,23 +59,29 @@ DepthwiseConvolution2dLayer::InferOutputShapes(const std::vector<TensorShape>& i
BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Convolutions will always have 4D input.");
- unsigned int inWidth = inputShape[3];
- unsigned int inHeight = inputShape[2];
+ DataLayoutIndexed dataLayoutIndex(m_Param.m_DataLayout);
+
+ unsigned int inWidth = inputShape[dataLayoutIndex.GetWidthIndex()];
+ unsigned int inHeight = inputShape[dataLayoutIndex.GetHeightIndex()];
unsigned int inBatchSize = inputShape[0];
- unsigned int filterWidth = filterShape[3];
+ unsigned int filterWidth = filterShape[dataLayoutIndex.GetWidthIndex()];
unsigned int readWidth = (inWidth + m_Param.m_PadLeft + m_Param.m_PadRight) - (filterWidth);
- unsigned int outWidth = 1+(readWidth / m_Param.m_StrideX);
+ unsigned int outWidth = 1 + (readWidth / m_Param.m_StrideX);
- unsigned int filterHeight = filterShape[2];
+ unsigned int filterHeight = filterShape[dataLayoutIndex.GetHeightIndex()];
unsigned int readHeight = (inHeight + m_Param.m_PadTop + m_Param.m_PadBottom) - (filterHeight);
- unsigned int outHeight = 1+(readHeight / m_Param.m_StrideY);
+ unsigned int outHeight = 1 + (readHeight / m_Param.m_StrideY);
unsigned int depthMultiplier = filterShape[0];
- unsigned int outChannels = filterShape[1]*depthMultiplier;
+ unsigned int outChannels = filterShape[dataLayoutIndex.GetChannelsIndex()] * depthMultiplier;
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 DepthwiseConvolution2dLayer::ValidateTensorShapesFromInputs()