commit ae7b832a6f5eda4b28577f57909111135a36dee9
author Narumol Prangnawarat <narumol.prangnawarat@arm.com> Fri Aug 02 15:08:59 2019 +0100
committer Narumol Prangnawarat <narumol.prangnawarat@arm.com> Fri Aug 02 15:26:58 2019 +0100
tree 3ea84622726dc3800c741d7be3cebb775eb20bb2
parent 87f65eab4abb65273ea11eb8ca876196ef82c6c5

IVGCVSW-3604 Fix channel shape calculation in TransposeConvolution2dLayer::InferOutputShapes

Signed-off-by: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Change-Id: I2e3d5922bb89c8f3b84ff5458fda981ff177c3ce

src/armnn/layers/TransposeConvolution2dLayer.cpp

diff --git a/src/armnn/layers/TransposeConvolution2dLayer.cpp b/src/armnn/layers/TransposeConvolution2dLayer.cpp
index 1a994e7..77a333d 100644
--- a/src/armnn/layers/TransposeConvolution2dLayer.cpp
+++ b/src/armnn/layers/TransposeConvolution2dLayer.cpp
@@ -67,7 +67,6 @@
     DataLayoutIndexed dataLayoutIndex(m_Param.m_DataLayout);
 
     const unsigned int batches  = inputShape[0];
-    const unsigned int channels = inputShape[dataLayoutIndex.GetChannelsIndex()];
 
     const unsigned int wInput = inputShape[dataLayoutIndex.GetWidthIndex()];
     const unsigned int hInput = inputShape[dataLayoutIndex.GetHeightIndex()];
@@ -84,6 +83,13 @@
     unsigned int wOutput = wPaddedOutput - (m_Param.m_PadLeft + m_Param.m_PadRight);
     unsigned int hOutput = hPaddedOutput - (m_Param.m_PadTop + m_Param.m_PadBottom);
 
+    unsigned int kernelElements = kernelShape[0] * kernelShape[dataLayoutIndex.GetChannelsIndex()];
+    unsigned int inputElements = batches * inputShape[dataLayoutIndex.GetChannelsIndex()];
+
+    BOOST_ASSERT_MSG(inputElements != 0, "Invalid number of input elements");
+    BOOST_ASSERT_MSG(kernelElements % inputElements == 0, "Invalid number of elements");
+    unsigned int channels =  kernelElements / inputElements;
+
     TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ?
          TensorShape( { batches, hOutput, wOutput, channels } ) :
          TensorShape( { batches, channels, hOutput, wOutput });

src/armnn/test/InferOutputTests.hpp

diff --git a/src/armnn/test/InferOutputTests.hpp b/src/armnn/test/InferOutputTests.hpp
index 2dd2ff0..c428a9d 100644
--- a/src/armnn/test/InferOutputTests.hpp
+++ b/src/armnn/test/InferOutputTests.hpp
@@ -406,7 +406,7 @@
     armnn::TensorShape filterShape(4, filterSize.data());
     shapes.push_back(filterShape);
 
-    const std::vector<unsigned int> expectedOutputSizes = {1, 2, 6, 6};
+    const std::vector<unsigned int> expectedOutputSizes = {1, 1, 6, 6};
     armnn::TensorShape expectedOutputShape(4, expectedOutputSizes.data());
 
     BOOST_CHECK(expectedOutputShape == transposeConvolution2dLayer->InferOutputShapes(shapes).at(0));