IVGCVSW-6382 Add Shape operator support to ONNX parser
Signed-off-by: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Change-Id: I3547effcbebf1ebc02d3b20f5db394a26991424d
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f64c8c4..28d63d3 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -751,6 +751,7 @@
src/armnnOnnxParser/test/ProtoxtFixture.cpp
src/armnnOnnxParser/test/Relu.cpp
src/armnnOnnxParser/test/Reshape.cpp
+ src/armnnOnnxParser/test/Shape.cpp
)
endif()
diff --git a/src/armnnOnnxParser/OnnxParser.cpp b/src/armnnOnnxParser/OnnxParser.cpp
index 49f0271..889c35f 100644
--- a/src/armnnOnnxParser/OnnxParser.cpp
+++ b/src/armnnOnnxParser/OnnxParser.cpp
@@ -426,7 +426,8 @@
{ "LeakyRelu", &OnnxParserImpl::ParseLeakyRelu },
{ "Conv", &OnnxParserImpl::ParseConv },
{ "Add", &OnnxParserImpl::ParseAdd },
- { "Flatten", &OnnxParserImpl::ParseFlatten},
+ { "Flatten", &OnnxParserImpl::ParseFlatten },
+ { "Shape", &OnnxParserImpl::ParseShape }
};
template<typename TypePair, typename Location>
@@ -1653,6 +1654,30 @@
AddPoolingLayer(node, desc);
}
+void OnnxParserImpl::ParseShape(const onnx::NodeProto& node)
+{
+ CHECK_VALID_SIZE(static_cast<size_t>(node.input_size()), 1);
+ CHECK_VALID_SIZE(static_cast<size_t>(node.output_size()), 1);
+
+ // Output must be INT64
+ CHECK_VALID_DATATYPE(node.name(), node.output(0),
+ m_TensorsInfo[node.output(0)].m_dtype,
+ onnx::TensorProto::INT64);
+
+ IConnectableLayer* layer = m_Network->AddShapeLayer(node.name().c_str());
+ ARMNN_ASSERT(layer != nullptr);
+
+ TensorShape inputShape = m_TensorsInfo[node.input(0)].m_info->GetShape();
+ auto outputInfo = ComputeOutputInfo({node.output(0)}, layer, {inputShape});
+ layer->GetOutputSlot(0).SetTensorInfo(outputInfo[0]);
+
+ // register the input connection slots for the layer, connections are made after all layers have been created
+ RegisterInputSlots(layer, {node.input(0)});
+
+ // register the output connection slots for the layer, connections are made after all layers have been created
+ RegisterOutputSlots(layer, {node.output(0)});
+}
+
void OnnxParserImpl::ParseReshape(const onnx::NodeProto& node)
{
CHECK_VALID_SIZE(static_cast<size_t>(node.input_size()), 2);
diff --git a/src/armnnOnnxParser/OnnxParser.hpp b/src/armnnOnnxParser/OnnxParser.hpp
index f618ff4..101e99f 100644
--- a/src/armnnOnnxParser/OnnxParser.hpp
+++ b/src/armnnOnnxParser/OnnxParser.hpp
@@ -117,6 +117,7 @@
void ParseFlatten(const onnx::NodeProto& node);
void ParseGlobalAveragePool(const onnx::NodeProto& node);
void ParseMaxPool(const onnx::NodeProto& nodeProto);
+ void ParseShape(const onnx::NodeProto& node);
void ParseReshape(const onnx::NodeProto& nodeProto);
void RegisterInputSlots(armnn::IConnectableLayer* layer, const std::vector<std::string>& tensorIndexes);
diff --git a/src/armnnOnnxParser/test/Shape.cpp b/src/armnnOnnxParser/test/Shape.cpp
new file mode 100644
index 0000000..b033b2d
--- /dev/null
+++ b/src/armnnOnnxParser/test/Shape.cpp
@@ -0,0 +1,144 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "armnnOnnxParser/IOnnxParser.hpp"
+#include "ParserPrototxtFixture.hpp"
+
+TEST_SUITE("OnnxParser_Shape")
+{
+struct ShapeMainFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+ ShapeMainFixture(const std::string& inputType,
+ const std::string& outputType,
+ const std::string& outputDim,
+ const std::vector<int>& inputShape)
+ {
+ m_Prototext = R"(
+ ir_version: 8
+ producer_name: "onnx-example"
+ graph {
+ node {
+ input: "Input"
+ output: "Output"
+ op_type: "Shape"
+ }
+ name: "shape-model"
+ input {
+ name: "Input"
+ type {
+ tensor_type {
+ elem_type: )" + inputType + R"(
+ shape {
+ )" + ConstructShapeString(inputShape) + R"(
+ }
+ }
+ }
+ }
+ output {
+ name: "Output"
+ type {
+ tensor_type {
+ elem_type: )" + outputType + R"(
+ shape {
+ dim {
+ dim_value: )" + outputDim + R"(
+ }
+ }
+ }
+ }
+ }
+ }
+ opset_import {
+ version: 10
+ })";
+ }
+ std::string ConstructShapeString(const std::vector<int>& shape)
+ {
+ std::string shapeStr;
+ for (int i : shape)
+ {
+ shapeStr = fmt::format("{} dim {{ dim_value: {} }}", shapeStr, i);
+ }
+ return shapeStr;
+ }
+};
+
+struct ShapeValidFloatFixture : ShapeMainFixture
+{
+ ShapeValidFloatFixture() : ShapeMainFixture("1", "7", "4", { 1, 3, 1, 5 }) {
+ Setup();
+ }
+};
+
+struct ShapeValidIntFixture : ShapeMainFixture
+{
+ ShapeValidIntFixture() : ShapeMainFixture("7", "7", "4", { 1, 3, 1, 5 }) {
+ Setup();
+ }
+};
+
+struct Shape3DFixture : ShapeMainFixture
+{
+ Shape3DFixture() : ShapeMainFixture("1", "7", "3", { 3, 2, 3 }) {
+ Setup();
+ }
+};
+
+struct Shape2DFixture : ShapeMainFixture
+{
+ Shape2DFixture() : ShapeMainFixture("1", "7", "2", { 2, 3 }) {
+ Setup();
+ }
+};
+
+struct Shape1DFixture : ShapeMainFixture
+{
+ Shape1DFixture() : ShapeMainFixture("1", "7", "1", { 5 }) {
+ Setup();
+ }
+};
+
+struct ShapeInvalidFixture : ShapeMainFixture
+{
+ ShapeInvalidFixture() : ShapeMainFixture("1", "1", "4", { 1, 3, 1, 5 }) {}
+};
+
+TEST_CASE_FIXTURE(ShapeValidFloatFixture, "FloatValidShapeTest")
+{
+ RunTest<2, int>({{"Input", { 0.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+ 4.0f, 3.0f, 2.0f, 1.0f, 0.0f,
+ 0.0f, 1.0f, 2.0f, 3.0f, 4.0f }}}, {{"Output", { 1, 3, 1, 5 }}});
+}
+
+TEST_CASE_FIXTURE(ShapeValidIntFixture, "IntValidShapeTest")
+{
+ RunTest<2, int>({{"Input", { 0, 1, 2, 3, 4,
+ 4, 3, 2, 1, 0,
+ 0, 1, 2, 3, 4 }}}, {{"Output", { 1, 3, 1, 5 }}});
+}
+
+TEST_CASE_FIXTURE(Shape3DFixture, "Shape3DTest")
+{
+ RunTest<2, int>({{"Input", { 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f,
+ 5.0f, 4.0f, 3.0f, 2.0f, 1.0f, 0.0f,
+ 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f }}}, {{"Output", { 3, 2, 3 }}});
+}
+
+TEST_CASE_FIXTURE(Shape2DFixture, "Shape2DTest")
+{
+ RunTest<2, int>({{"Input", { 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f }}}, {{"Output", { 2, 3 }}});
+}
+
+TEST_CASE_FIXTURE(Shape1DFixture, "Shape1DTest")
+{
+ RunTest<2, int>({{"Input", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f }}}, {{"Output", { 5 }}});
+}
+
+TEST_CASE_FIXTURE(ShapeInvalidFixture, "IncorrectOutputDataShapeTest")
+{
+ CHECK_THROWS_AS(Setup(), armnn::ParseException);
+}
+
+}