IVGCVSW-6450 Add Support of Models with Dynamic Batch Tensor to ONNX parser

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

1 files changed, 26 insertions, 0 deletions

diff --git a/tests/InferenceModel.hpp b/tests/InferenceModel.hpp
index 02511965d9..cf3aae137e 100644
--- a/tests/InferenceModel.hpp
+++ b/tests/InferenceModel.hpp
@@ -312,6 +312,32 @@ public:
 
         armnn::INetworkPtr network{nullptr, [](armnn::INetwork *){}};
 
+        std::map<std::string, armnn::TensorShape> inputShapes;
+        if (!params.m_InputShapes.empty())
+        {
+            const size_t numInputShapes   = params.m_InputShapes.size();
+            const size_t numInputBindings = params.m_InputBindings.size();
+            if (numInputShapes < numInputBindings)
+            {
+                throw armnn::Exception(fmt::format(
+                    "Not every input has its tensor shape specified: expected={0}, got={1}",
+                    numInputBindings, numInputShapes));
+            }
+
+            for (size_t i = 0; i < numInputShapes; i++)
+            {
+                inputShapes[params.m_InputBindings[i]] = params.m_InputShapes[i];
+            }
+
+            {
+                ARMNN_SCOPED_HEAP_PROFILING("Parsing");
+                network = (params.m_IsModelBinary ?
+                    parser->CreateNetworkFromBinaryFile(modelPath.c_str(), inputShapes) :
+                    parser->CreateNetworkFromTextFile(modelPath.c_str(), inputShapes));
+            }
+        }
+
+        else
         {
             ARMNN_SCOPED_HEAP_PROFILING("Parsing");
             network = (params.m_IsModelBinary ?