IVGCVSW-2083 - Add DataLayout parameter to calculate the expected output shape in ValidateTensorShapesFromInputs

 * Convolution2dLayer
 * DepthwiseConvolution2dLayer
 * Pooling2dLayer
 * ResizeBilinearLayer
 * Unittests for ValidateTensorShapesFromInputs

Change-Id: I057caf8a90d822175a7dd7271f960b65c6154bb4

5 files changed, 231 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()
diff --git a/src/armnn/test/OptimizerTests.cpp b/src/armnn/test/OptimizerTests.cpp
index 9d351babd8..30df6eb0e6 100644
--- a/src/armnn/test/OptimizerTests.cpp
+++ b/src/armnn/test/OptimizerTests.cpp
@@ -771,4 +771,194 @@ BOOST_AUTO_TEST_CASE(Fp32NetworkToFp16OptimizationTest)
                              &IsLayerOfType<armnn::OutputLayer>));
 }
 
+void CreateConvolution2dGraph(Graph &graph, const unsigned int* inputShape,
+                              const unsigned int* weightsShape, const unsigned int* outputShape,
+                              DataLayout dataLayout = DataLayout::NCHW)
+{
+    armnn::TensorInfo inputInfo(4, inputShape, DataType::Float32);
+    armnn::TensorInfo outputInfo(4, outputShape, DataType::Float32);
+
+    std::vector<float> weightsVector(90);
+    armnn::ConstTensor weights(armnn::TensorInfo(4, weightsShape, armnn::DataType::Float32), weightsVector);
+
+    Convolution2dDescriptor desc;
+    desc.m_BiasEnabled = false;
+    desc.m_StrideX = 1;
+    desc.m_StrideY = 1;
+    desc.m_DataLayout = dataLayout;
+
+    Layer* input = graph.AddLayer<InputLayer>(0, "input");
+    input->GetOutputSlot().SetTensorInfo(inputInfo);
+
+    Convolution2dLayer* layer = graph.AddLayer<Convolution2dLayer>(desc, "conv2d");
+    layer->m_Weight = std::make_unique<armnn::ScopedCpuTensorHandle>(weights);
+    layer->GetOutputSlot().SetTensorInfo(outputInfo);
+
+    Layer* output = graph.AddLayer<OutputLayer>(0, "output");
+    input->GetOutputSlot().Connect(layer->GetInputSlot(0));
+    layer->GetOutputSlot().Connect(output->GetInputSlot(0));
+}
+
+BOOST_AUTO_TEST_CASE(Conv2dValidateTensorShapesFromInputs)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 3, 8, 16 };
+    const unsigned int weightsShape[] = { 2, 3, 5, 3 };
+    const unsigned int outputShape[] = { 1, 2, 4, 14 };
+    CreateConvolution2dGraph(graph, inputShape, weightsShape, outputShape);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+BOOST_AUTO_TEST_CASE(Conv2dValidateTensorShapesFromInputsNhwc)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 8, 16, 3 };
+    const unsigned int weightsShape[] = { 2, 5, 3, 3 };
+    const unsigned int outputShape[] = { 1, 4, 14, 2 };
+    CreateConvolution2dGraph(graph, inputShape, weightsShape, outputShape, DataLayout::NHWC);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+void CreateDepthwiseConvolution2dGraph(Graph &graph, const unsigned int* inputShape,
+                              const unsigned int* weightsShape, const unsigned int* outputShape,
+                              DataLayout dataLayout = DataLayout::NCHW)
+{
+    armnn::TensorInfo inputInfo(4, inputShape, DataType::Float32);
+    armnn::TensorInfo outputInfo(4, outputShape, DataType::Float32);
+
+    std::vector<float> weightsVector(18);
+    armnn::ConstTensor weights(armnn::TensorInfo(4, weightsShape, armnn::DataType::Float32), weightsVector);
+
+    DepthwiseConvolution2dDescriptor desc;
+    desc.m_BiasEnabled = false;
+    desc.m_StrideX = 1;
+    desc.m_StrideY = 1;
+    desc.m_DataLayout = dataLayout;
+
+    Layer* input = graph.AddLayer<InputLayer>(0, "input");
+    input->GetOutputSlot().SetTensorInfo(inputInfo);
+
+    DepthwiseConvolution2dLayer* layer = graph.AddLayer<DepthwiseConvolution2dLayer>(desc, "depthwiseConv2d");
+    layer->m_Weight = std::make_unique<armnn::ScopedCpuTensorHandle>(weights);
+    layer->GetOutputSlot().SetTensorInfo(outputInfo);
+
+    Layer* output = graph.AddLayer<OutputLayer>(0, "output");
+    input->GetOutputSlot().Connect(layer->GetInputSlot(0));
+    layer->GetOutputSlot().Connect(output->GetInputSlot(0));
+}
+
+BOOST_AUTO_TEST_CASE(DepthwiseConv2dValidateTensorShapesFromInputs)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 2, 3, 3 };
+    const unsigned int weightsShape[] = { 1, 2, 3, 3 };
+    const unsigned int outputShape[] = { 1, 2, 1, 1 };
+    CreateDepthwiseConvolution2dGraph(graph, inputShape, weightsShape, outputShape);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+BOOST_AUTO_TEST_CASE(DepthwiseConv2dValidateTensorShapesFromInputsNhwc)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 3, 3, 2 };
+    const unsigned int weightsShape[] = { 1, 3, 3, 2 };
+    const unsigned int outputShape[] = { 1, 1, 1, 2 };
+    CreateDepthwiseConvolution2dGraph(graph, inputShape, weightsShape, outputShape, DataLayout::NHWC);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+void CreatePooling2dGraph(Graph &graph, const unsigned int* inputShape,  const unsigned int* outputShape,
+                          DataLayout dataLayout = DataLayout::NCHW)
+{
+    armnn::TensorInfo inputInfo(4, inputShape, DataType::Float32);
+    armnn::TensorInfo outputInfo(4, outputShape, DataType::Float32);
+
+    Pooling2dDescriptor desc;
+    desc.m_PoolType = armnn::PoolingAlgorithm::Average;
+    desc.m_PoolWidth = desc.m_PoolHeight = 100;
+    desc.m_StrideX = desc.m_StrideY = 5;
+    desc.m_PadLeft = 50;
+    desc.m_PadRight = 50;
+    desc.m_PadTop = 50;
+    desc.m_PadBottom = 50;
+    desc.m_PaddingMethod = armnn::PaddingMethod::Exclude;
+    desc.m_DataLayout = dataLayout;
+
+    Layer* input = graph.AddLayer<InputLayer>(0, "input");
+    input->GetOutputSlot().SetTensorInfo(inputInfo);
+
+    Pooling2dLayer* layer = graph.AddLayer<Pooling2dLayer>(desc, "pooling2d");
+    layer->GetOutputSlot().SetTensorInfo(outputInfo);
+
+    Layer* output = graph.AddLayer<OutputLayer>(0, "output");
+    input->GetOutputSlot().Connect(layer->GetInputSlot(0));
+    layer->GetOutputSlot().Connect(output->GetInputSlot(0));
+}
+
+BOOST_AUTO_TEST_CASE(Pooling2dValidateTensorShapesFromInputs)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 5, 3, 52, 60 };
+    const unsigned int outputShape[] = { 5, 3, 11, 13 };
+    CreatePooling2dGraph(graph, inputShape, outputShape, DataLayout::NCHW);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+BOOST_AUTO_TEST_CASE(Pooling2dValidateTensorShapesFromInputsNhwc)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 5, 52, 60, 3 };
+    const unsigned int outputShape[] = { 5, 11, 13, 3 };
+    CreatePooling2dGraph(graph, inputShape, outputShape, DataLayout::NHWC);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+void CreateResizeBilinearGraph(Graph &graph, const unsigned int* inputShape,  const unsigned int* outputShape,
+                               DataLayout dataLayout = DataLayout::NCHW)
+{
+    armnn::TensorInfo inputInfo(4, inputShape, DataType::Float32);
+    armnn::TensorInfo outputInfo(4, outputShape, DataType::Float32);
+
+    ResizeBilinearDescriptor desc;
+    desc.m_TargetHeight = 3;
+    desc.m_TargetWidth = 4;
+    desc.m_DataLayout = dataLayout;
+
+    Layer* input = graph.AddLayer<InputLayer>(0, "input");
+    input->GetOutputSlot().SetTensorInfo(inputInfo);
+
+    ResizeBilinearLayer* layer = graph.AddLayer<ResizeBilinearLayer>(desc, "resizeBilinear");
+    layer->GetOutputSlot().SetTensorInfo(outputInfo);
+
+    Layer* output = graph.AddLayer<OutputLayer>(0, "output");
+    input->GetOutputSlot().Connect(layer->GetInputSlot(0));
+    layer->GetOutputSlot().Connect(output->GetInputSlot(0));
+}
+
+BOOST_AUTO_TEST_CASE(ResizeBilinearValidateTensorShapesFromInputs)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 2, 4, 5 };
+    const unsigned int outputShape[] = { 1, 2, 3, 4 };
+    CreateResizeBilinearGraph(graph, inputShape, outputShape);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
+BOOST_AUTO_TEST_CASE(ResizeBilinearValidateTensorShapesFromInputsNhwc)
+{
+    Graph graph;
+    const unsigned int inputShape[] = { 1, 4, 5, 2 };
+    const unsigned int outputShape[] = { 1, 3, 4, 2 };
+    CreateResizeBilinearGraph(graph, inputShape, outputShape, DataLayout::NHWC);
+
+    BOOST_CHECK_NO_THROW(graph.InferTensorInfos());
+}
+
 BOOST_AUTO_TEST_SUITE_END()