commit fcf6fd562f87595c814d8acbec04194421018c32
author Tee Jung <tee.ty.jung@openedges.com> Fri Nov 01 05:27:28 2019 +0000
committer Matteo Martincigh <matteo.martincigh@arm.com> Mon Nov 04 09:11:08 2019 +0000
tree 0d719705d6d05388025457ee9592ed35aa6309c0
parent d94efa8a7c6476fb50b3434723fda22859c236ad

Fix crash issue

* armnnOnnxParser makes tensorInfo from graph->value_info
  but PyTorch does not add weight/bias tensor information to graph->value_info
  so tensorInfo of const tensor should be extracted from graph->initializer

Signed-off-by: Jung Tae-young tee.ty.jung@openedges.com
Change-Id: Ib2656dd25abc522012cf413e843fe03949cb2eb0

src/armnnOnnxParser/OnnxParser.cpp

diff --git a/src/armnnOnnxParser/OnnxParser.cpp b/src/armnnOnnxParser/OnnxParser.cpp
index 6f83dc5..9d374ae 100644
--- a/src/armnnOnnxParser/OnnxParser.cpp
+++ b/src/armnnOnnxParser/OnnxParser.cpp
@@ -187,16 +187,10 @@
     return attribValue;
 }
 
-armnn::TensorInfo ToTensorInfo(const onnx::ValueInfoProto& info)
+armnn::TensorInfo ToTensorInfo(const std::string& name, std::vector<unsigned int>& shape, int data_type)
 {
-  const onnx::TensorShapeProto onnxShape = info.type().tensor_type().shape();
-  std::vector<unsigned int> shapeDims;
-  for (int i = 0; i < onnxShape.dim_size(); ++i)
-  {
-      shapeDims.push_back(CHECKED_NON_NEGATIVE(CHECKED_INT32(onnxShape.dim(i).dim_value())));
-  }
   DataType type;
-  switch(info.type().tensor_type().elem_type())
+  switch(data_type)
   {
       case onnx::TensorProto::FLOAT:
       {
@@ -216,13 +210,35 @@
                   boost::format("'%1%' is not a currently supported datatype for tensor %2%."
                                 " Supported dataTypes are FLOAT, INT32 and INT64.  %3%") %
                                 onnx::TensorProto::DataType_Name(
-                                    static_cast<onnx::TensorProto::DataType>(info.type().tensor_type().elem_type())) %
-                                info.name() %
+                                    static_cast<onnx::TensorProto::DataType>(data_type)) %
+                                name %
                                 CHECK_LOCATION().AsString() ));
       }
-
   }
-  return TensorInfo(TensorShape(static_cast<unsigned int>(shapeDims.size()), shapeDims.data()), type);
+  return TensorInfo(TensorShape(static_cast<unsigned int>(shape.size()), shape.data()), type);
+}
+
+armnn::TensorInfo ToTensorInfo(const onnx::ValueInfoProto& info)
+{
+  const onnx::TensorShapeProto onnxShape = info.type().tensor_type().shape();
+  std::vector<unsigned int> shapeDims;
+  for (int i = 0; i < onnxShape.dim_size(); ++i)
+  {
+      shapeDims.push_back(CHECKED_NON_NEGATIVE(CHECKED_INT32(onnxShape.dim(i).dim_value())));
+  }
+
+  return ToTensorInfo(info.name(), shapeDims, info.type().tensor_type().elem_type());
+}
+
+armnn::TensorInfo ToTensorInfo(const onnx::TensorProto& tensor)
+{
+  std::vector<unsigned int> shapeDims;
+  for (auto dim: tensor.dims())
+  {
+      shapeDims.push_back(CHECKED_NON_NEGATIVE(CHECKED_INT32(dim)));
+  }
+
+  return ToTensorInfo(tensor.name(), shapeDims, tensor.data_type());
 }
 
 std::string TensorInfoAsString(const TensorInfo& info,
@@ -580,6 +596,9 @@
     for (auto tensor : m_Graph->initializer())
     {
         m_TensorsInfo[tensor.name()].m_tensor = std::make_unique<const onnx::TensorProto>(tensor);
+        m_TensorsInfo[tensor.name()].m_info = std::make_unique<TensorInfo>(ToTensorInfo(tensor));
+        m_TensorsInfo[tensor.name()].m_dtype =
+            static_cast<onnx::TensorProto::DataType>(tensor.data_type());
     }
 
     SetupInputLayers();