From f77cab57b3eca1425384d4d5bfe44d76fc7023b9 Mon Sep 17 00:00:00 2001 From: Teresa Charlin Date: Thu, 1 Jun 2023 16:15:13 +0100 Subject: IVGCVSW-7785 Extend support for 3D tensors BATCH_TO_SPACE and SPACE_TO_BATCH in CpuRef * Both layers were assuming 4D tensors, now 3D is supported too. * Remove some unnecessary includes * Add Unit Tests Signed-off-by: Teresa Charlin Change-Id: I7bdd11e4936a27cd97ec65fd915e6ccaa1494cff --- src/armnn/layers/BatchToSpaceNdLayer.cpp | 46 +++++++++----------------------- src/armnn/layers/SpaceToBatchNdLayer.cpp | 31 +++++++++------------ 2 files changed, 25 insertions(+), 52 deletions(-) (limited to 'src/armnn/layers') diff --git a/src/armnn/layers/BatchToSpaceNdLayer.cpp b/src/armnn/layers/BatchToSpaceNdLayer.cpp index f022c525a8..b760b5661c 100644 --- a/src/armnn/layers/BatchToSpaceNdLayer.cpp +++ b/src/armnn/layers/BatchToSpaceNdLayer.cpp @@ -1,18 +1,11 @@ // -// Copyright © 2017 Arm Ltd and Contributors. All rights reserved. +// Copyright © 2017,2023 Arm Ltd and Contributors. All rights reserved. // SPDX-License-Identifier: MIT // #include "BatchToSpaceNdLayer.hpp" #include "LayerCloneBase.hpp" -#include "LayerWithParameters.hpp" -#include "BatchToSpaceNdLayer.hpp" - -#include - -#include -#include #include #include @@ -59,8 +52,6 @@ void BatchToSpaceNdLayer::ValidateTensorShapesFromInputs() std::vector BatchToSpaceNdLayer::InferOutputShapes(const std::vector& inputShapes) const { - ARMNN_ASSERT(inputShapes.size() == 1); - const TensorShape& inputShape = inputShapes[0]; TensorShape outputShape(inputShape); @@ -68,29 +59,18 @@ std::vector BatchToSpaceNdLayer::InferOutputShapes(const std::vecto m_Param.m_BlockShape.end(), 1U, std::multiplies<>()); - - ARMNN_ASSERT(inputShape[0] % accumulatedBlockShape == 0); - - outputShape[0] = inputShape[0] / accumulatedBlockShape; - - DataLayoutIndexed dimensionIndices = m_Param.m_DataLayout; - unsigned int heightIndex = dimensionIndices.GetHeightIndex(); - unsigned int widthIndex = dimensionIndices.GetWidthIndex(); - - unsigned int heightCrop = m_Param.m_Crops[0].first + m_Param.m_Crops[0].second; - unsigned int widthCrop = m_Param.m_Crops[1].first + m_Param.m_Crops[1].second; - - unsigned int outputHeight = inputShape[heightIndex] * m_Param.m_BlockShape[0]; - unsigned int outputWidth = inputShape[widthIndex] * m_Param.m_BlockShape[1]; - - ARMNN_ASSERT_MSG(heightCrop <= outputHeight, - "BatchToSpaceLayer: Overall height crop should be less than or equal to the uncropped output height."); - - ARMNN_ASSERT_MSG(widthCrop <= outputWidth, - "BatchToSpaceLayer: Overall width crop should be less than or equal to the uncropped output width."); - - outputShape[heightIndex] = outputHeight - heightCrop; - outputShape[widthIndex] = outputWidth - widthCrop; + outputShape[0] = (inputShape[0] / accumulatedBlockShape) < 1 ? 1 : (inputShape[0] / accumulatedBlockShape) ; + + // In a 4D tensor, there will be 2 spatialDimensions (H and W), and the for loop will run twice. + // In a 3D tensor, there will be 1 spatialDimensions, and the for loop will run once. + unsigned int firstSpatialDimension = m_Param.m_DataLayout == DataLayout::NCHW ? 2 : 1; + for (unsigned int i = 0; i < m_Param.m_BlockShape.size(); ++i) + { + unsigned int spatialDimension = firstSpatialDimension + i; + unsigned int cropSize = m_Param.m_Crops[i].first + m_Param.m_Crops[i].second; + unsigned int outputSize = inputShape[spatialDimension] * m_Param.m_BlockShape[i]; + outputShape[spatialDimension] = outputSize - cropSize; + } return std::vector({ outputShape }); } diff --git a/src/armnn/layers/SpaceToBatchNdLayer.cpp b/src/armnn/layers/SpaceToBatchNdLayer.cpp index 151b6a5301..a758617e2e 100644 --- a/src/armnn/layers/SpaceToBatchNdLayer.cpp +++ b/src/armnn/layers/SpaceToBatchNdLayer.cpp @@ -1,15 +1,11 @@ // -// Copyright © 2017 Arm Ltd and Contributors. All rights reserved. +// Copyright © 2017,2023 Arm Ltd and Contributors. All rights reserved. // SPDX-License-Identifier: MIT // #include "SpaceToBatchNdLayer.hpp" #include "LayerCloneBase.hpp" -#include - -#include - #include #include @@ -42,9 +38,7 @@ SpaceToBatchNdLayer* SpaceToBatchNdLayer::Clone(Graph& graph) const std::vector SpaceToBatchNdLayer::InferOutputShapes(const std::vector& inputShapes) const { - ARMNN_ASSERT(inputShapes.size() == 1); - - TensorShape inputShape = inputShapes[0]; + const TensorShape inputShape = inputShapes[0]; TensorShape outputShape(inputShape); outputShape[0] = inputShape[0] * std::accumulate(m_Param.m_BlockShape.begin(), @@ -52,17 +46,16 @@ std::vector SpaceToBatchNdLayer::InferOutputShapes(const std::vecto 1U, std::multiplies<>()); - DataLayoutIndexed dimensionIndices = m_Param.m_DataLayout; - unsigned int heightIndex = dimensionIndices.GetHeightIndex(); - unsigned int widthIndex = dimensionIndices.GetWidthIndex(); - - std::pair heightPad = m_Param.m_PadList[0]; - std::pair widthPad = m_Param.m_PadList[1]; - - outputShape[heightIndex] = - (inputShape[heightIndex] + heightPad.first + heightPad.second) / m_Param.m_BlockShape[0]; - outputShape[widthIndex] = - (inputShape[widthIndex] + widthPad.first + widthPad.second) / m_Param.m_BlockShape[1]; + // In a 4D tensor, there will be 2 spatialDimensions (H and W), and the for loop will run twice. + // In a 3D tensor, there will be 1 spatialDimensions, and the for loop will run once. + unsigned int firstSpatialDimension = m_Param.m_DataLayout == DataLayout::NCHW ? 2 : 1; + for (unsigned int i = 0; i < m_Param.m_BlockShape.size(); ++i) + { + unsigned int spatialDimension = firstSpatialDimension + i; + outputShape[spatialDimension] = + (inputShape[spatialDimension] + m_Param.m_PadList[i].first + m_Param.m_PadList[i].second) + / m_Param.m_BlockShape[i]; + } return std::vector({ outputShape }); } -- cgit v1.2.1