IVGCVSW-7104: BatchMatMul Serializer/Deserializer Support
* Updated FlatBuffers schema for BatchMatMul layer type
* Added Serializer and Deserializer implementations for BatchMatMul
* Added unit tests for BatchMatMul serialization and deserialization
* Updated CMakeLists and docs
Signed-off-by: Samuel Yap <samuel.yap@arm.com>
Change-Id: Iad63afbd036a3eb648683eb7416a475561aa20cb
diff --git a/CMakeLists.txt b/CMakeLists.txt
index c63d8fc..4e4818d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -775,6 +775,7 @@
src/armnnDeserializer/test/DeserializeActivation.cpp
src/armnnDeserializer/test/DeserializeAdd.cpp
src/armnnDeserializer/test/DeserializeArgMinMax.cpp
+ src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
src/armnnDeserializer/test/DeserializeBatchToSpaceNd.cpp
src/armnnDeserializer/test/DeserializeBatchNormalization.cpp
src/armnnDeserializer/test/DeserializeCast.cpp
diff --git a/docs/05_02_deserializer_serializer.dox b/docs/05_02_deserializer_serializer.dox
index 6cfaf29..c36e010 100644
--- a/docs/05_02_deserializer_serializer.dox
+++ b/docs/05_02_deserializer_serializer.dox
@@ -22,6 +22,7 @@
- Activation
- Addition
- ArgMinMax
+- BatchMatMul
- BatchToSpaceNd
- BatchNormalization
- Cast
@@ -114,6 +115,7 @@
- Activation
- Addition
- ArgMinMax
+- BatchMatMul
- BatchToSpaceNd
- BatchNormalization
- Cast
diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp
index a405cb9..702b060 100644
--- a/src/armnnDeserializer/Deserializer.cpp
+++ b/src/armnnDeserializer/Deserializer.cpp
@@ -214,6 +214,7 @@
m_ParserFunctions[Layer_ActivationLayer] = &DeserializerImpl::ParseActivation;
m_ParserFunctions[Layer_AdditionLayer] = &DeserializerImpl::ParseAdd;
m_ParserFunctions[Layer_ArgMinMaxLayer] = &DeserializerImpl::ParseArgMinMax;
+ m_ParserFunctions[Layer_BatchMatMulLayer] = &DeserializerImpl::ParseBatchMatMul;
m_ParserFunctions[Layer_BatchToSpaceNdLayer] = &DeserializerImpl::ParseBatchToSpaceNd;
m_ParserFunctions[Layer_BatchNormalizationLayer] = &DeserializerImpl::ParseBatchNormalization;
m_ParserFunctions[Layer_CastLayer] = &DeserializerImpl::ParseCast;
@@ -292,6 +293,8 @@
return graphPtr->layers()->Get(layerIndex)->layer_as_AdditionLayer()->base();
case Layer::Layer_ArgMinMaxLayer:
return graphPtr->layers()->Get(layerIndex)->layer_as_ArgMinMaxLayer()->base();
+ case Layer::Layer_BatchMatMulLayer:
+ return graphPtr->layers()->Get(layerIndex)->layer_as_BatchMatMulLayer()->base();
case Layer::Layer_BatchToSpaceNdLayer:
return graphPtr->layers()->Get(layerIndex)->layer_as_BatchToSpaceNdLayer()->base();
case Layer::Layer_BatchNormalizationLayer:
@@ -1258,6 +1261,37 @@
RegisterOutputSlots(graph, layerIndex, layer);
}
+void IDeserializer::DeserializerImpl::ParseBatchMatMul(GraphPtr graph, unsigned int layerIndex)
+{
+ CHECK_LAYERS(graph, 0, layerIndex);
+
+ auto inputs = GetInputs(graph, layerIndex);
+ CHECK_LOCATION();
+ CHECK_VALID_SIZE(inputs.size(), 2);
+
+ auto outputs = GetOutputs(graph, layerIndex);
+ CHECK_VALID_SIZE(outputs.size(), 1);
+
+ auto serializerLayer = graph->layers()->Get(layerIndex)->layer_as_BatchMatMulLayer();
+ auto serializerDescriptor = serializerLayer->descriptor();
+
+ armnn::BatchMatMulDescriptor descriptor(serializerDescriptor->transposeX(),
+ serializerDescriptor->transposeY(),
+ serializerDescriptor->adjointX(),
+ serializerDescriptor->adjointY(),
+ ToDataLayout(serializerDescriptor->dataLayoutX()),
+ ToDataLayout(serializerDescriptor->dataLayoutY()));
+
+ auto layerName = GetLayerName(graph, layerIndex);
+ IConnectableLayer* layer = m_Network->AddBatchMatMulLayer(descriptor, layerName.c_str());
+
+ armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0]);
+ layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo);
+
+ RegisterInputSlots(graph, layerIndex, layer);
+ RegisterOutputSlots(graph, layerIndex, layer);
+}
+
void IDeserializer::DeserializerImpl::ParseBatchToSpaceNd(GraphPtr graph, unsigned int layerIndex)
{
CHECK_LAYERS(graph, 0, layerIndex);
diff --git a/src/armnnDeserializer/Deserializer.hpp b/src/armnnDeserializer/Deserializer.hpp
index 277c09a..bd01a35 100644
--- a/src/armnnDeserializer/Deserializer.hpp
+++ b/src/armnnDeserializer/Deserializer.hpp
@@ -88,6 +88,7 @@
void ParseActivation(GraphPtr graph, unsigned int layerIndex);
void ParseAdd(GraphPtr graph, unsigned int layerIndex);
void ParseArgMinMax(GraphPtr graph, unsigned int layerIndex);
+ void ParseBatchMatMul(GraphPtr graph, unsigned int layerIndex);
void ParseBatchToSpaceNd(GraphPtr graph, unsigned int layerIndex);
void ParseBatchNormalization(GraphPtr graph, unsigned int layerIndex);
void ParseCast(GraphPtr graph, unsigned int layerIndex);
diff --git a/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp b/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
new file mode 100644
index 0000000..40f93ce
--- /dev/null
+++ b/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
@@ -0,0 +1,213 @@
+//
+// Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "ParserFlatbuffersSerializeFixture.hpp"
+#include <armnnDeserializer/IDeserializer.hpp>
+
+#include <doctest/doctest.h>
+
+#include <string>
+
+TEST_SUITE("Deserializer_BatchMatMul")
+{
+struct BatchMatMulFixture : public ParserFlatbuffersSerializeFixture
+{
+ explicit BatchMatMulFixture(const std::string& inputXShape,
+ const std::string& inputYShape,
+ const std::string& outputShape,
+ const std::string& dataType)
+ {
+ m_JsonString = R"(
+ {
+ inputIds:[
+ 0,
+ 1
+ ],
+ outputIds:[
+ 3
+ ],
+ layers:[
+ {
+ layer_type:"InputLayer",
+ layer:{
+ base:{
+ layerBindingId:0,
+ base:{
+ index:0,
+ layerName:"InputXLayer",
+ layerType:"Input",
+ inputSlots:[
+ {
+ index:0,
+ connection:{
+ sourceLayerIndex:0,
+ outputSlotIndex:0
+ },
+
+ }
+ ],
+ outputSlots:[
+ {
+ index:0,
+ tensorInfo:{
+ dimensions:)" + inputXShape + R"(,
+ dataType:)" + dataType + R"(
+ },
+
+ }
+ ],
+
+ },
+
+ }
+ },
+
+ },
+ {
+ layer_type:"InputLayer",
+ layer:{
+ base:{
+ layerBindingId:1,
+ base:{
+ index:1,
+ layerName:"InputYLayer",
+ layerType:"Input",
+ inputSlots:[
+ {
+ index:0,
+ connection:{
+ sourceLayerIndex:0,
+ outputSlotIndex:0
+ },
+
+ }
+ ],
+ outputSlots:[
+ {
+ index:0,
+ tensorInfo:{
+ dimensions:)" + inputYShape + R"(,
+ dataType:)" + dataType + R"(
+ },
+
+ }
+ ],
+
+ },
+
+ }
+ },
+
+ },
+ {
+ layer_type:"BatchMatMulLayer",
+ layer:{
+ base:{
+ index:2,
+ layerName:"BatchMatMulLayer",
+ layerType:"BatchMatMul",
+ inputSlots:[
+ {
+ index:0,
+ connection:{
+ sourceLayerIndex:0,
+ outputSlotIndex:0
+ },
+
+ },
+ {
+ index:1,
+ connection:{
+ sourceLayerIndex:1,
+ outputSlotIndex:0
+ },
+
+ }
+ ],
+ outputSlots:[
+ {
+ index:0,
+ tensorInfo:{
+ dimensions:)" + outputShape + R"(,
+ dataType:)" + dataType + R"(
+ },
+
+ }
+ ],
+
+ },
+ descriptor:{
+ transposeX:false,
+ transposeY:false,
+ adjointX:false,
+ adjointY:false,
+ dataLayoutX:NHWC,
+ dataLayoutY:NHWC
+ }
+ },
+
+ },
+ {
+ layer_type:"OutputLayer",
+ layer:{
+ base:{
+ layerBindingId:0,
+ base:{
+ index:3,
+ layerName:"OutputLayer",
+ layerType:"Output",
+ inputSlots:[
+ {
+ index:0,
+ connection:{
+ sourceLayerIndex:2,
+ outputSlotIndex:0
+ },
+
+ }
+ ],
+ outputSlots:[
+ {
+ index:0,
+ tensorInfo:{
+ dimensions:)" + outputShape + R"(,
+ dataType:)" + dataType + R"(
+ },
+
+ }
+ ],
+
+ }
+ }
+ },
+
+ }
+ ]
+ }
+ )";
+ Setup();
+ }
+};
+
+struct SimpleBatchMatMulFixture : BatchMatMulFixture
+{
+ SimpleBatchMatMulFixture()
+ : BatchMatMulFixture("[ 1, 2, 2, 1 ]",
+ "[ 1, 2, 2, 1 ]",
+ "[ 1, 2, 2, 1 ]",
+ "Float32")
+ {}
+};
+
+TEST_CASE_FIXTURE(SimpleBatchMatMulFixture, "SimpleBatchMatMulTest")
+{
+ RunTest<4, armnn::DataType::Float32>(
+ 0,
+ {{"InputXLayer", { 1.0f, 2.0f, 3.0f, 4.0f }},
+ {"InputYLayer", { 5.0f, 6.0f, 7.0f, 8.0f }}},
+ {{"OutputLayer", { 19.0f, 22.0f, 43.0f, 50.0f }}});
+}
+
+}
\ No newline at end of file
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs
index f301fce..2dbfd85 100644
--- a/src/armnnSerializer/ArmnnSchema.fbs
+++ b/src/armnnSerializer/ArmnnSchema.fbs
@@ -182,6 +182,7 @@
Convolution3d = 65,
Pooling3d = 66,
GatherNd = 67,
+ BatchMatMul = 68,
}
// Base layer table to be used as part of other layers
@@ -1009,6 +1010,20 @@
inputParams:LstmInputParams;
}
+table BatchMatMulDescriptor {
+ transposeX:bool = false;
+ transposeY:bool = false;
+ adjointX:bool = false;
+ adjointY:bool = false;
+ dataLayoutX:DataLayout = NCHW;
+ dataLayoutY:DataLayout = NCHW;
+}
+
+table BatchMatMulLayer {
+ base:LayerBase;
+ descriptor:BatchMatMulDescriptor;
+}
+
union Layer {
ActivationLayer,
AdditionLayer,
@@ -1078,6 +1093,7 @@
Convolution3dLayer,
Pooling3dLayer,
GatherNdLayer,
+ BatchMatMulLayer,
}
table AnyLayer {
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index 488dac6..c9a3022 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -218,6 +218,33 @@
CreateAnyLayer(flatBufferLayer.o, serializer::Layer::Layer_ArgMinMaxLayer);
}
+void SerializerStrategy::SerializeBatchMatMulLayer(const armnn::IConnectableLayer* layer,
+ const armnn::BatchMatMulDescriptor& descriptor,
+ const char* name)
+{
+ IgnoreUnused(name);
+
+ // Create FlatBuffer BaseLayer
+ auto flatBufferBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_BatchMatMul);
+
+ // Create the FlatBuffer BatchMatMulDescriptor
+ auto flatBufferDescriptor = CreateBatchMatMulDescriptor(m_flatBufferBuilder,
+ descriptor.m_TransposeX,
+ descriptor.m_TransposeY,
+ descriptor.m_AdjointX,
+ descriptor.m_AdjointY,
+ GetFlatBufferDataLayout(descriptor.m_DataLayoutX),
+ GetFlatBufferDataLayout(descriptor.m_DataLayoutY));
+
+ // Create the FlatBuffer BatchMatMulLayer
+ auto flatBufferBatchMatMulLayer = CreateBatchMatMulLayer(m_flatBufferBuilder,
+ flatBufferBaseLayer,
+ flatBufferDescriptor);
+
+ // Add the AnyLayer to the FlatBufferLayers
+ CreateAnyLayer(flatBufferBatchMatMulLayer.o, serializer::Layer::Layer_BatchMatMulLayer);
+}
+
// Build FlatBuffer for BatchToSpaceNd Layer
void SerializerStrategy::SerializeBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer,
const armnn::BatchToSpaceNdDescriptor& descriptor,
@@ -1971,6 +1998,15 @@
SerializeArgMinMaxLayer(layer, layerDescriptor, name);
break;
}
+ case armnn::LayerType::BatchMatMul:
+ {
+ const armnn::BatchMatMulDescriptor& layerDescriptor =
+ static_cast<const armnn::BatchMatMulDescriptor&>(descriptor);
+ SerializeBatchMatMulLayer(layer,
+ layerDescriptor,
+ name);
+ break;
+ }
case armnn::LayerType::BatchNormalization :
{
const armnn::BatchNormalizationDescriptor& layerDescriptor =
diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp
index 216f4dc..60fed4f 100644
--- a/src/armnnSerializer/Serializer.hpp
+++ b/src/armnnSerializer/Serializer.hpp
@@ -113,6 +113,10 @@
const armnn::ArgMinMaxDescriptor& argMinMaxDescriptor,
const char* name = nullptr);
+ void SerializeBatchMatMulLayer(const armnn::IConnectableLayer* layer,
+ const armnn::BatchMatMulDescriptor& descriptor,
+ const char* name = nullptr);
+
void SerializeBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer,
const armnn::BatchToSpaceNdDescriptor& descriptor,
const char* name = nullptr);
diff --git a/src/armnnSerializer/test/SerializerTests.cpp b/src/armnnSerializer/test/SerializerTests.cpp
index 3c00fc4..a568bf1 100644
--- a/src/armnnSerializer/test/SerializerTests.cpp
+++ b/src/armnnSerializer/test/SerializerTests.cpp
@@ -92,6 +92,47 @@
SerializeArgMinMaxTest(armnn::DataType::Signed64);
}
+TEST_CASE("SerializeBatchMatMul")
+{
+ const std::string layerName("batchMatMul");
+ const armnn::TensorInfo inputXInfo({2, 3, 4, 5}, armnn::DataType::Float32);
+ const armnn::TensorInfo inputYInfo({2, 4, 3, 5}, armnn::DataType::Float32);
+
+ const armnn::TensorInfo outputInfo({2, 3, 3, 5}, armnn::DataType::Float32);
+
+ armnn::BatchMatMulDescriptor descriptor(false,
+ false,
+ false,
+ false,
+ armnn::DataLayout::NHWC,
+ armnn::DataLayout::NHWC);
+
+ armnn::INetworkPtr network = armnn::INetwork::Create();
+ armnn::IConnectableLayer* const inputXLayer = network->AddInputLayer(0);
+ armnn::IConnectableLayer* const inputYLayer = network->AddInputLayer(1);
+
+ armnn::IConnectableLayer* const batchMatMulLayer =
+ network->AddBatchMatMulLayer(descriptor, layerName.c_str());
+ armnn::IConnectableLayer* const outputLayer = network->AddOutputLayer(0);
+
+ inputXLayer->GetOutputSlot(0).Connect(batchMatMulLayer->GetInputSlot(0));
+ inputYLayer->GetOutputSlot(0).Connect(batchMatMulLayer->GetInputSlot(1));
+ batchMatMulLayer->GetOutputSlot(0).Connect(outputLayer->GetInputSlot(0));
+
+ inputXLayer->GetOutputSlot(0).SetTensorInfo(inputXInfo);
+ inputYLayer->GetOutputSlot(0).SetTensorInfo(inputYInfo);
+ batchMatMulLayer->GetOutputSlot(0).SetTensorInfo(outputInfo);
+
+ armnn::INetworkPtr deserializedNetwork = DeserializeNetwork(SerializeNetwork(*network));
+ CHECK(deserializedNetwork);
+
+ LayerVerifierBaseWithDescriptor<armnn::BatchMatMulDescriptor> verifier(layerName,
+ {inputXInfo, inputYInfo},
+ {outputInfo},
+ descriptor);
+ deserializedNetwork->ExecuteStrategy(verifier);
+}
+
TEST_CASE("SerializeBatchNormalization")
{
const std::string layerName("batchNormalization");