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

7 files changed, 371 insertions, 9 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8fd71239eb..c84749679a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -768,6 +768,7 @@ if(BUILD_UNIT_TESTS)
             src/armnnOnnxParser/test/Gather.cpp
             src/armnnOnnxParser/test/Gemm.cpp
             src/armnnOnnxParser/test/GetInputsOutputs.cpp
+            src/armnnOnnxParser/test/LoadScopeDynamicTensor.cpp
             src/armnnOnnxParser/test/OnnxParserTestUtils.cpp
             src/armnnOnnxParser/test/OnnxParserTestUtils.hpp
             src/armnnOnnxParser/test/Pooling.cpp
diff --git a/include/armnnOnnxParser/IOnnxParser.hpp b/include/armnnOnnxParser/IOnnxParser.hpp
index f9d692738d..ba7fc83f93 100644
--- a/include/armnnOnnxParser/IOnnxParser.hpp
+++ b/include/armnnOnnxParser/IOnnxParser.hpp
@@ -36,6 +36,19 @@ public:
     /// Create the network directly from protobuf text in a string. Useful for debugging/testing
     armnn::INetworkPtr CreateNetworkFromString(const std::string& protoText);
 
+    /// Create the network from a protobuf binary file on disk, with inputShapes specified
+    armnn::INetworkPtr CreateNetworkFromBinaryFile(const char* graphFile,
+                                                   const std::map<std::string, armnn::TensorShape>& inputShapes);
+
+    /// Create the network from a protobuf text file on disk, with inputShapes specified
+    armnn::INetworkPtr CreateNetworkFromTextFile(const char* graphFile,
+                                                 const std::map<std::string, armnn::TensorShape>& inputShapes);
+
+     /// Create the network directly from protobuf text in a string, with inputShapes specified.
+     /// Useful for debugging/testing
+    armnn::INetworkPtr CreateNetworkFromString(const std::string& protoText,
+                                               const std::map<std::string, armnn::TensorShape>& inputShapes);
+
     /// Retrieve binding info (layer id and tensor info) for the network input identified by the given layer name
     BindingPointInfo GetNetworkInputBindingInfo(const std::string& name) const;
 
diff --git a/src/armnnOnnxParser/OnnxParser.cpp b/src/armnnOnnxParser/OnnxParser.cpp
index 3588975897..eb24bb5425 100644
--- a/src/armnnOnnxParser/OnnxParser.cpp
+++ b/src/armnnOnnxParser/OnnxParser.cpp
@@ -60,6 +60,25 @@ armnn::INetworkPtr IOnnxParser::CreateNetworkFromString(const std::string& proto
     return pOnnxParserImpl->CreateNetworkFromString(protoText);
 }
 
+armnn::INetworkPtr IOnnxParser::CreateNetworkFromBinaryFile(
+    const char* graphFile,
+    const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    return pOnnxParserImpl->CreateNetworkFromBinaryFile(graphFile, inputShapes);
+}
+
+armnn::INetworkPtr IOnnxParser::CreateNetworkFromTextFile(const char* graphFile,
+                                                          const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    return pOnnxParserImpl->CreateNetworkFromTextFile(graphFile, inputShapes);
+}
+
+armnn::INetworkPtr IOnnxParser::CreateNetworkFromString(const std::string& protoText,
+                                                        const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    return pOnnxParserImpl->CreateNetworkFromString(protoText, inputShapes);
+}
+
 BindingPointInfo IOnnxParser::GetNetworkInputBindingInfo(const std::string& name) const
 {
     return pOnnxParserImpl->GetNetworkInputBindingInfo(name);
@@ -287,12 +306,18 @@ armnn::TensorInfo ToTensorInfo(const std::string& name, std::vector<unsigned int
         }
     }
 
-    // To avoid crashes by trivial tensors
+    // Scalar Tensor
     if (shape.empty())
     {
         return TensorInfo(TensorShape(Dimensionality::Scalar), type);
     }
 
+    // Dynamic Tensor
+    if(std::find(shape.begin(), shape.end(), 0) != shape.end())
+    {
+        return TensorInfo(TensorShape(Dimensionality::NotSpecified), type);
+    }
+
     return TensorInfo(TensorShape(static_cast<unsigned int>(shape.size()), shape.data()), type);
 }
 
@@ -469,7 +494,9 @@ std::vector<TensorInfo> OnnxParserImpl::ComputeOutputInfo(std::vector<std::strin
                                    outNames.end(),
                                    [this](std::string name)
                                    {
-                                       return (m_TensorsInfo.count(name) == 0 || m_TensorsInfo[name].m_info == nullptr);
+                                       return (m_TensorsInfo.count(name) == 0 || m_TensorsInfo[name].m_info == nullptr
+                                       || m_TensorsInfo[name].m_info->GetShape().GetDimensionality() ==
+                                          Dimensionality::NotSpecified);
                                    });
     std::vector<TensorInfo> outInfo;
     //if the output info(s) are not here, we need to compute them
@@ -521,6 +548,8 @@ void OnnxParserImpl::ResetParser()
 {
     m_Network = armnn::INetworkPtr(nullptr, nullptr);
     m_Graph = nullptr;
+    m_InputInfos.clear();
+    m_OutputInfos.clear();
 }
 
 void OnnxParserImpl::Cleanup()
@@ -529,6 +558,7 @@ void OnnxParserImpl::Cleanup()
     m_TensorsInfo.clear();
     m_OutputsMap.clear();
     m_OutputsFusedAndUsed.clear();
+    m_InputShapes.clear();
 }
 
 template<typename T>
@@ -692,6 +722,14 @@ INetworkPtr OnnxParserImpl::CreateNetworkFromTextFile(const char* graphFile)
     return CreateNetworkFromModel(*modelProto);
 }
 
+INetworkPtr OnnxParserImpl::CreateNetworkFromTextFile(const char* graphFile,
+                                                      const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    ResetParser();
+    m_InputShapes = inputShapes;
+    ModelPtr modelProto = LoadModelFromTextFile(graphFile);
+    return CreateNetworkFromModel(*modelProto);
+}
 
 ModelPtr OnnxParserImpl::LoadModelFromBinaryFile(const char* graphFile)
 {
@@ -728,6 +766,15 @@ INetworkPtr OnnxParserImpl::CreateNetworkFromBinaryFile(const char* graphFile)
     return CreateNetworkFromModel(*modelProto);
 }
 
+INetworkPtr OnnxParserImpl::CreateNetworkFromBinaryFile(const char* graphFile,
+                                                        const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    ResetParser();
+    m_InputShapes = inputShapes;
+    ModelPtr modelProto = LoadModelFromBinaryFile(graphFile);
+    return CreateNetworkFromModel(*modelProto);
+}
+
 ModelPtr OnnxParserImpl::LoadModelFromString(const std::string& protoText)
 {
     if (protoText == "")
@@ -754,6 +801,15 @@ INetworkPtr OnnxParserImpl::CreateNetworkFromString(const std::string& protoText
     return CreateNetworkFromModel(*modelProto);
 }
 
+INetworkPtr OnnxParserImpl::CreateNetworkFromString(const std::string& protoText,
+                                                    const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    ResetParser();
+    m_InputShapes = inputShapes;
+    ModelPtr modelProto = LoadModelFromString(protoText);
+    return CreateNetworkFromModel(*modelProto);
+}
+
 INetworkPtr OnnxParserImpl::CreateNetworkFromModel(onnx::ModelProto& model)
 {
     m_Network = INetwork::Create();
@@ -843,6 +899,13 @@ void OnnxParserImpl::LoadGraph()
             }
         }
     }
+
+    // Get output info.
+    for(int outputIndex = 0; outputIndex < m_Graph->output_size(); ++outputIndex)
+    {
+        auto output = m_Graph->output(outputIndex);
+        m_OutputInfos[output.name()] = *m_TensorsInfo[output.name()].m_info;
+    }
 }
 
 void OnnxParserImpl::SetupInfo(const google::protobuf::RepeatedPtrField<onnx::ValueInfoProto >* list)
@@ -2172,13 +2235,31 @@ void OnnxParserImpl::SetupInputLayers()
     for(int inputIndex = 0; inputIndex < m_Graph->input_size(); ++inputIndex)
     {
         auto input = m_Graph->input(inputIndex);
-        if (! m_TensorsInfo[input.name()].isConstant())
+        if (!m_TensorsInfo[input.name()].isConstant())
         {
             IConnectableLayer* layer =
-              m_Network->AddInputLayer(static_cast<armnn::LayerBindingId>(inputIndex), input.name().c_str());
-            auto tensorInfo = ToTensorInfo(input);
+                m_Network->AddInputLayer(static_cast<armnn::LayerBindingId>(inputIndex), input.name().c_str());
+            TensorInfo tensorInfo = *m_TensorsInfo[input.name()].m_info;
+            if (tensorInfo.GetShape().GetDimensionality() == Dimensionality::NotSpecified)
+            {
+                if (m_InputShapes.find(input.name()) == m_InputShapes.end())
+                {
+                    throw ParseException(fmt::format("The parser does not support dynamic tensor, "
+                                                     "please specify input shape for {}. {}",
+                                                     input.name(),
+                                                     CHECK_LOCATION().AsString()));
+                }
+                else
+                {
+                    tensorInfo.SetShape(m_InputShapes[input.name()]);
+                    m_TensorsInfo[input.name()].m_info = std::make_unique<TensorInfo>(tensorInfo);
+                }
+
+            }
             layer->GetOutputSlot(0).SetTensorInfo(tensorInfo);
 
+            m_InputInfos[input.name()] = tensorInfo;
+
             RegisterOutputSlots(layer,{ input.name() });
         }
     }
@@ -2211,7 +2292,7 @@ void OnnxParserImpl::RegisterInputSlot(IConnectableLayer* layer,
 
     if (it == m_TensorConnections.end())
     {
-        //First time seing this tensor, we need to map it
+        //First time seeing this tensor, we need to map it
         m_TensorConnections[tensorId] = TensorSlots();
     }
     m_TensorConnections[tensorId].inputSlots.push_back(slot);
@@ -2238,7 +2319,7 @@ void OnnxParserImpl::RegisterInputSlots(IConnectableLayer* layer, const std::vec
 
         if (it == m_TensorConnections.end())
         {
-            //First time seing this tensor, we need to map it
+            // First time seing this tensor, we need to map it
             m_TensorConnections[tensorId] = TensorSlots();
         }
         m_TensorConnections[tensorId].inputSlots.push_back(slot);
@@ -2282,6 +2363,7 @@ void OnnxParserImpl::RegisterOutputSlots(IConnectableLayer* layer, const std::ve
         }
         tensorSlots.outputSlot = slot;
     }
+
 }
 
 BindingPointInfo OnnxParserImpl::GetNetworkInputBindingInfo(const std::string& name) const
@@ -2291,7 +2373,12 @@ BindingPointInfo OnnxParserImpl::GetNetworkInputBindingInfo(const std::string& n
         auto input = m_Graph->input(i);
         if(input.name() == name)
         {
-            return std::make_pair(static_cast<armnn::LayerBindingId>(i), ToTensorInfo(input));
+            auto it = m_InputInfos.find(name);
+
+            if (it != m_InputInfos.end())
+            {
+                return std::make_pair(static_cast<armnn::LayerBindingId>(i), it->second);
+            }
         }
     }
     throw InvalidArgumentException(fmt::format("The input layer '{}' does not exist {}",
@@ -2305,7 +2392,12 @@ BindingPointInfo OnnxParserImpl::GetNetworkOutputBindingInfo(const std::string&
         auto output = m_Graph->output(i);
         if(output.name() == name)
         {
-            return std::make_pair(static_cast<armnn::LayerBindingId>(i), ToTensorInfo(output));
+            auto it = m_OutputInfos.find(name);
+
+            if (it != m_OutputInfos.end())
+            {
+                return std::make_pair(static_cast<armnn::LayerBindingId>(i), it->second);
+            }
         }
     }
     throw InvalidArgumentException(fmt::format("The output layer '{}' does not exist {}",
diff --git a/src/armnnOnnxParser/OnnxParser.hpp b/src/armnnOnnxParser/OnnxParser.hpp
index ec19006be7..a0fa8af31b 100644
--- a/src/armnnOnnxParser/OnnxParser.hpp
+++ b/src/armnnOnnxParser/OnnxParser.hpp
@@ -34,12 +34,25 @@ public:
     /// Create the network from a protobuf binary file on disk
     armnn::INetworkPtr CreateNetworkFromBinaryFile(const char* graphFile);
 
+    /// Create the network from a protobuf binary file on disk, with inputShapes specified
+    armnn::INetworkPtr CreateNetworkFromBinaryFile(const char* graphFile,
+                                                   const std::map<std::string, armnn::TensorShape>& inputShapes);
+
     /// Create the network from a protobuf text file on disk
     armnn::INetworkPtr CreateNetworkFromTextFile(const char* graphFile);
 
+    /// Create the network from a protobuf text file on disk, with inputShapes specified
+    armnn::INetworkPtr CreateNetworkFromTextFile(const char* graphFile,
+                                                 const std::map<std::string, armnn::TensorShape>& inputShapes);
+
     /// Create the network directly from protobuf text in a string. Useful for debugging/testing
     armnn::INetworkPtr CreateNetworkFromString(const std::string& protoText);
 
+     /// Create the network directly from protobuf text in a string, with inputShapes specified.
+     /// Useful for debugging/testing
+    armnn::INetworkPtr CreateNetworkFromString(const std::string& protoText,
+                                               const std::map<std::string, armnn::TensorShape>& inputShapes);
+
     /// Retrieve binding info (layer id and tensor info) for the network input identified by the given layer name
     BindingPointInfo GetNetworkInputBindingInfo(const std::string& name) const;
 
@@ -203,5 +216,11 @@ private:
 
     std::vector<UsageSummary> m_OutputsFusedAndUsed;
 
+    std::map<std::string, armnn::TensorShape> m_InputShapes;
+
+    std::unordered_map<std::string, armnn::TensorInfo> m_InputInfos;
+
+    std::unordered_map<std::string, armnn::TensorInfo> m_OutputInfos;
+
 };
 }
diff --git a/src/armnnOnnxParser/test/LoadScopeDynamicTensor.cpp b/src/armnnOnnxParser/test/LoadScopeDynamicTensor.cpp
new file mode 100644
index 0000000000..7d29d807af
--- /dev/null
+++ b/src/armnnOnnxParser/test/LoadScopeDynamicTensor.cpp
@@ -0,0 +1,184 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "armnnOnnxParser/IOnnxParser.hpp"
+#include "ParserPrototxtFixture.hpp"
+
+TEST_SUITE("OnnxParser_LoadScopeDynamicTensor")
+{
+
+struct DynamicBatchTensorFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    DynamicBatchTensorFixture()
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: 1
+                            shape {
+                              dim {
+                                dim_value: 0
+                              }
+                              dim {
+                                dim_value: 1
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                      input {
+                        name: "Weight"
+                        type {
+                          tensor_type {
+                            elem_type: 1
+                            shape {
+                              dim {
+                                dim_value: 1
+                              }
+                              dim {
+                                dim_value: 1
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                      initializer {
+                          dims: 1
+                          dims: 1
+                          dims: 3
+                          dims: 3
+                          data_type: 1
+                          float_data: 2
+                          float_data: 1
+                          float_data: 0
+                          float_data: 6
+                          float_data: 2
+                          float_data: 1
+                          float_data: 4
+                          float_data: 1
+                          float_data: 2
+                          name: "Weight"
+                        }
+                      node {
+                         input: "Input"
+                         input: "Weight"
+                         output: "Output"
+                         name: "Convolution"
+                         op_type: "Conv"
+                         attribute {
+                           name: "kernel_shape"
+                           ints: 3
+                           ints: 3
+                           type: INTS
+                         }
+                         attribute {
+                           name: "strides"
+                           ints: 1
+                           ints: 1
+                           type: INTS
+                         }
+                         attribute {
+                           name: "auto_pad"
+                           s: "VALID"
+                           type: STRING
+                         }
+                         attribute {
+                           name: "group"
+                           i: 1
+                           type: INT
+                         }
+                         attribute {
+                           name: "dilations"
+                           ints: 1
+                           ints: 1
+                           type: INTS
+                         }
+                         doc_string: ""
+                         domain: ""
+                       }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 0
+                                   }
+                                   dim {
+                                       dim_value: 1
+                                   }
+                                   dim {
+                                       dim_value: 1
+                                   }
+                                   dim {
+                                       dim_value: 1
+                                   }
+                               }
+                            }
+                        }
+                        }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+TEST_CASE_FIXTURE(DynamicBatchTensorFixture, "DynamicBatchTensorTest")
+{
+    Setup({{"Input", armnn::TensorShape({1, 1, 3, 3})}});
+    RunTest<4>({{"Input", {1.0, 2.0, 3.0,
+                           4.0, 5.0, 6.0,
+                           7.0, 8.0, 9.0}}},
+              {{"Output", {1.0 * 2 + 2.0 * 1 + 3.0 * 0 +
+                           4.0 * 6 + 5.0 * 2 + 6.0 * 1 +
+                           7.0 * 4 + 8.0 * 1 + 9.0 * 2}}});
+}
+
+TEST_CASE_FIXTURE(DynamicBatchTensorFixture, "TensorShapeNotSpecifiedTest")
+{
+    CHECK_THROWS_AS(Setup(), armnn::ParseException);
+}
+
+TEST_CASE_FIXTURE(DynamicBatchTensorFixture, "IncorrectInputNameTest")
+{
+    CHECK_THROWS_AS(Setup({{"Incorrect", armnn::TensorShape({1, 1, 3, 3})}}), armnn::ParseException);
+}
+
+TEST_CASE_FIXTURE(DynamicBatchTensorFixture, "IncorrectBatchTensorTest")
+{
+    Setup({{"Input", armnn::TensorShape({2, 1, 3, 3}) }});
+    CHECK_THROWS_AS(RunTest<4>({{"Input", { 1.0, 2.0, 3.0,
+                                            4.0, 5.0, 6.0,
+                                            7.0, 8.0, 9.0 }}},
+                               {{"Output", {1.0 * 2 + 2.0 * 1 + 3.0 * 0 +
+                                            4.0 * 6 + 5.0 * 2 + 6.0 * 1 +
+                                            7.0 * 4 + 8.0 * 1 + 9.0 * 2 }}}), armnn::Exception);
+
+}
+
+}
diff --git a/src/armnnUtils/ParserPrototxtFixture.hpp b/src/armnnUtils/ParserPrototxtFixture.hpp
index 08ac3aeb9b..3c659d3fd6 100644
--- a/src/armnnUtils/ParserPrototxtFixture.hpp
+++ b/src/armnnUtils/ParserPrototxtFixture.hpp
@@ -42,6 +42,7 @@ struct ParserPrototxtFixture
                                       const std::string& outputName);
     void Setup(const std::map<std::string, armnn::TensorShape>& inputShapes,
         const std::vector<std::string>& requestedOutputs);
+    void Setup(const std::map<std::string, armnn::TensorShape>& inputShapes);
     void Setup();
     armnn::IOptimizedNetworkPtr SetupOptimizedNetwork(
         const std::map<std::string,armnn::TensorShape>& inputShapes,
@@ -136,6 +137,23 @@ void ParserPrototxtFixture<TParser>::Setup(const std::map<std::string, armnn::Te
 }
 
 template<typename TParser>
+void ParserPrototxtFixture<TParser>::Setup(const std::map<std::string, armnn::TensorShape>& inputShapes)
+{
+    std::string errorMessage;
+
+    armnn::INetworkPtr network =
+        m_Parser->CreateNetworkFromString(m_Prototext.c_str(), inputShapes);
+    auto optimized = Optimize(*network, { armnn::Compute::CpuRef }, m_Runtime->GetDeviceSpec());
+    armnn::Status ret = m_Runtime->LoadNetwork(m_NetworkIdentifier, move(optimized), errorMessage);
+    if (ret != armnn::Status::Success)
+    {
+        throw armnn::Exception(fmt::format("LoadNetwork failed with error: '{0}' {1}",
+                                           errorMessage,
+                                           CHECK_LOCATION().AsString()));
+    }
+}
+
+template<typename TParser>
 void ParserPrototxtFixture<TParser>::Setup()
 {
     std::string errorMessage;
@@ -191,6 +209,15 @@ void ParserPrototxtFixture<TParser>::RunTest(const std::map<std::string, std::ve
     {
         armnn::BindingPointInfo bindingInfo = m_Parser->GetNetworkInputBindingInfo(it.first);
         inputTensors.push_back({ bindingInfo.first, armnn::ConstTensor(bindingInfo.second, it.second.data()) });
+        if (bindingInfo.second.GetNumElements() != it.second.size())
+        {
+            throw armnn::Exception(fmt::format("Input tensor {0} is expected to have {1} elements. "
+                                               "{2} elements supplied. {3}",
+                                               it.first,
+                                               bindingInfo.second.GetNumElements(),
+                                               it.second.size(),
+                                               CHECK_LOCATION().AsString()));
+        }
     }
 
     // Allocates storage for the output tensors to be written to and sets up the armnn output tensors.
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 ?