From c0ed7baa8c05c4710034dfd179fadd31b716a46f Mon Sep 17 00:00:00 2001
From: SiCong Li <sicong.li@arm.com>
Date: Fri, 21 Jun 2019 16:02:40 +0100
Subject: MLCE-103 Remove hardcoded output shape in ModelAccuracyTool

We can obtain the output tensor shape from the model provided by
the user.

Signed-off-by: SiCong Li <sicong.li@arm.com>
Change-Id: I5074734315174c1b5dc8eea1eff18a4a1c566f2a
---
 tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

(limited to 'tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp')

diff --git a/tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp b/tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp
index bb0d824e0e..85241e889c 100644
--- a/tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp
+++ b/tests/ModelAccuracyTool-Armnn/ModelAccuracyTool-Armnn.cpp
@@ -194,6 +194,9 @@ int main(int argc, char* argv[])
                 inputTensorDataLayout == armnn::DataLayout::NCHW ? inputTensorShape[3] : inputTensorShape[2];
             const unsigned int inputTensorHeight =
                 inputTensorDataLayout == armnn::DataLayout::NCHW ? inputTensorShape[2] : inputTensorShape[1];
+            // Get output tensor info
+            const unsigned int outputNumElements = model.GetOutputSize();
+
             const unsigned int batchSize = 1;
             // Get normalisation parameters
             SupportedFrontend modelFrontend;
@@ -232,7 +235,7 @@ int main(int argc, char* argv[])
                             normParams,
                             batchSize,
                             inputTensorDataLayout));
-                        outputDataContainers = {vector<int>(1001)};
+                        outputDataContainers = { vector<int>(outputNumElements) };
                         break;
                     case armnn::DataType::QuantisedAsymm8:
                         inputDataContainers.push_back(
@@ -241,7 +244,7 @@ int main(int argc, char* argv[])
                             normParams,
                             batchSize,
                             inputTensorDataLayout));
-                        outputDataContainers = {vector<uint8_t>(1001)};
+                        outputDataContainers = { vector<uint8_t>(outputNumElements) };
                         break;
                     case armnn::DataType::Float32:
                     default:
@@ -251,7 +254,7 @@ int main(int argc, char* argv[])
                             normParams,
                             batchSize,
                             inputTensorDataLayout));
-                        outputDataContainers = {vector<float>(1001)};
+                        outputDataContainers = { vector<float>(outputNumElements) };
                         break;
                 }
 
-- 
cgit v1.2.1