diff options
Diffstat (limited to 'src/armnn/layers')
-rw-r--r-- | src/armnn/layers/Convolution2dLayer.cpp | 20 | ||||
-rw-r--r-- | src/armnn/layers/DepthwiseConvolution2dLayer.cpp | 22 | ||||
-rw-r--r-- | src/armnn/layers/Pooling2dLayer.cpp | 13 | ||||
-rw-r--r-- | src/armnn/layers/ResizeBilinearLayer.cpp | 8 |
4 files changed, 41 insertions, 22 deletions
diff --git a/src/armnn/layers/Convolution2dLayer.cpp b/src/armnn/layers/Convolution2dLayer.cpp index d4b67cca3f..d611aedc06 100644 --- a/src/armnn/layers/Convolution2dLayer.cpp +++ b/src/armnn/layers/Convolution2dLayer.cpp @@ -58,22 +58,28 @@ std::vector<TensorShape> Convolution2dLayer::InferOutputShapes(const std::vector // If we support multiple batch dimensions in the future, then this assert will need to change. 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 outChannels = filterShape[0]; 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 Convolution2dLayer::ValidateTensorShapesFromInputs() 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() 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() diff --git a/src/armnn/layers/ResizeBilinearLayer.cpp b/src/armnn/layers/ResizeBilinearLayer.cpp index 9f0608d11c..fda93da99f 100644 --- a/src/armnn/layers/ResizeBilinearLayer.cpp +++ b/src/armnn/layers/ResizeBilinearLayer.cpp @@ -37,10 +37,14 @@ std::vector<TensorShape> ResizeBilinearLayer::InferOutputShapes(const std::vecto unsigned int outWidth = m_Param.m_TargetWidth; unsigned int outHeight = m_Param.m_TargetHeight; - unsigned int outChannels = inputShape[1]; + unsigned int outChannels = inputShape[m_Param.m_DataLayout.GetChannelsIndex()]; unsigned int outBatch = inputShape[0]; - return std::vector<TensorShape>({ TensorShape({outBatch, outChannels, outHeight, outWidth}) }); + TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ? + TensorShape( { outBatch, outHeight, outWidth, outChannels } ) : + TensorShape( { outBatch, outChannels, outHeight, outWidth }); + + return std::vector<TensorShape>({ tensorShape }); } void ResizeBilinearLayer::ValidateTensorShapesFromInputs() |