IVGCVSW-3886 Add serialization support for DepthToSpace

Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Change-Id: Id80c1bcbf50715f07a92dad08d554518a283cc28

4 files changed, 80 insertions, 3 deletions

diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs
index 6b07510a21..6450d16ca8 100644
--- a/src/armnnSerializer/ArmnnSchema.fbs
+++ b/src/armnnSerializer/ArmnnSchema.fbs
@@ -140,7 +140,8 @@ enum LayerType : uint {
     QuantizedLstm = 45,
     Abs = 46,
     ArgMinMax = 47,
-    Slice = 48
+    Slice = 48,
+    DepthToSpace = 49
 }
 
 // Base layer table to be used as part of other layers
@@ -213,6 +214,16 @@ table Convolution2dDescriptor {
     dataLayout:DataLayout = NCHW;
 }
 
+table DepthToSpaceLayer {
+    base:LayerBase;
+    descriptor:DepthToSpaceDescriptor;
+}
+
+table DepthToSpaceDescriptor {
+    blockSize:uint;
+    dataLayout:DataLayout;
+}
+
 table DivisionLayer {
     base:LayerBase;
 }
@@ -722,7 +733,8 @@ union Layer {
     StackLayer,
     AbsLayer,
     ArgMinMaxLayer,
-    SliceLayer
+    SliceLayer,
+    DepthToSpaceLayer
 }
 
 table AnyLayer {
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index f07805cb38..99ba7e3c13 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -303,7 +303,14 @@ void SerializerVisitor::VisitDepthToSpaceLayer(const armnn::IConnectableLayer* l
                                                const armnn::DepthToSpaceDescriptor& descriptor,
                                                const char* name)
 {
-    throw UnimplementedException("SerializerVisitor::VisitDepthToSpaceLayer is not implemented");
+    auto fbBaseLayer  = CreateLayerBase(layer, serializer::LayerType::LayerType_DepthToSpace);
+    auto fbDescriptor = CreateDepthToSpaceDescriptor(m_flatBufferBuilder,
+                                                     descriptor.m_BlockSize,
+                                                     GetFlatBufferDataLayout(descriptor.m_DataLayout));
+
+    auto fbLayer = serializer::CreateDepthToSpaceLayer(m_flatBufferBuilder, fbBaseLayer, fbDescriptor);
+
+    CreateAnyLayer(fbLayer.o, serializer::Layer::Layer_DepthToSpaceLayer);
 }
 
 void SerializerVisitor::VisitDepthwiseConvolution2dLayer(const armnn::IConnectableLayer* layer,
diff --git a/src/armnnSerializer/SerializerSupport.md b/src/armnnSerializer/SerializerSupport.md
index 136998c3f2..cc9e59a762 100644
--- a/src/armnnSerializer/SerializerSupport.md
+++ b/src/armnnSerializer/SerializerSupport.md
@@ -15,6 +15,7 @@ The Arm NN SDK Serializer currently supports the following layers:
 * Concat
 * Constant
 * Convolution2d
+* DepthToSpace
 * DepthwiseConvolution2d
 * Dequantize
 * DetectionPostProcess
diff --git a/src/armnnSerializer/test/SerializerTests.cpp b/src/armnnSerializer/test/SerializerTests.cpp
index d4ad6236b4..dfa98acbcd 100644
--- a/src/armnnSerializer/test/SerializerTests.cpp
+++ b/src/armnnSerializer/test/SerializerTests.cpp
@@ -658,6 +658,63 @@ BOOST_AUTO_TEST_CASE(SerializeConvolution2d)
     deserializedNetwork->Accept(verifier);
 }
 
+BOOST_AUTO_TEST_CASE(SerializeDepthToSpace)
+{
+    class DepthToSpaceLayerVerifier : public LayerVerifierBase
+    {
+    public:
+        DepthToSpaceLayerVerifier(const std::string& layerName,
+                                  const std::vector<armnn::TensorInfo>& inputInfos,
+                                  const std::vector<armnn::TensorInfo>& outputInfos,
+                                  const armnn::DepthToSpaceDescriptor& descriptor)
+            : LayerVerifierBase(layerName, inputInfos, outputInfos)
+            , m_Descriptor(descriptor) {}
+
+        void VisitDepthToSpaceLayer(const armnn::IConnectableLayer* layer,
+                                    const armnn::DepthToSpaceDescriptor& descriptor,
+                                    const char* name) override
+        {
+            VerifyNameAndConnections(layer, name);
+            VerifyDescriptor(descriptor);
+        }
+
+    private:
+        void VerifyDescriptor(const armnn::DepthToSpaceDescriptor& descriptor)
+        {
+            BOOST_TEST(descriptor.m_BlockSize == m_Descriptor.m_BlockSize);
+            BOOST_TEST(GetDataLayoutName(descriptor.m_DataLayout) == GetDataLayoutName(m_Descriptor.m_DataLayout));
+        }
+
+        armnn::DepthToSpaceDescriptor m_Descriptor;
+    };
+
+    const std::string layerName("depthToSpace");
+
+    const armnn::TensorInfo inputInfo ({ 1,  8, 4, 12 }, armnn::DataType::Float32);
+    const armnn::TensorInfo outputInfo({ 1, 16, 8,  3 }, armnn::DataType::Float32);
+
+    armnn::DepthToSpaceDescriptor desc;
+    desc.m_BlockSize  = 2;
+    desc.m_DataLayout = armnn::DataLayout::NHWC;
+
+    armnn::INetworkPtr network = armnn::INetwork::Create();
+    armnn::IConnectableLayer* const inputLayer        = network->AddInputLayer(0);
+    armnn::IConnectableLayer* const depthToSpaceLayer = network->AddDepthToSpaceLayer(desc, layerName.c_str());
+    armnn::IConnectableLayer* const outputLayer       = network->AddOutputLayer(0);
+
+    inputLayer->GetOutputSlot(0).Connect(depthToSpaceLayer->GetInputSlot(0));
+    depthToSpaceLayer->GetOutputSlot(0).Connect(outputLayer->GetInputSlot(0));
+
+    inputLayer->GetOutputSlot(0).SetTensorInfo(inputInfo);
+    depthToSpaceLayer->GetOutputSlot(0).SetTensorInfo(outputInfo);
+
+    armnn::INetworkPtr deserializedNetwork = DeserializeNetwork(SerializeNetwork(*network));
+    BOOST_CHECK(deserializedNetwork);
+
+    DepthToSpaceLayerVerifier verifier(layerName, {inputInfo}, {outputInfo}, desc);
+    deserializedNetwork->Accept(verifier);
+}
+
 BOOST_AUTO_TEST_CASE(SerializeDepthwiseConvolution2d)
 {
     class DepthwiseConvolution2dLayerVerifier : public LayerVerifierBase