IVGCVSW-3886 Add serialization support for DepthToSpace

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

7 files changed, 111 insertions, 3 deletions

diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp
index 1c4d86cf34..d887357ef1 100644
--- a/src/armnnDeserializer/Deserializer.cpp
+++ b/src/armnnDeserializer/Deserializer.cpp
@@ -194,6 +194,7 @@ m_ParserFunctions(Layer_MAX+1, &Deserializer::ParseUnsupportedLayer)
     m_ParserFunctions[Layer_ConcatLayer]                 = &Deserializer::ParseConcat;
     m_ParserFunctions[Layer_ConstantLayer]               = &Deserializer::ParseConstant;
     m_ParserFunctions[Layer_Convolution2dLayer]          = &Deserializer::ParseConvolution2d;
+    m_ParserFunctions[Layer_DepthToSpaceLayer]           = &Deserializer::ParseDepthToSpace;
     m_ParserFunctions[Layer_DepthwiseConvolution2dLayer] = &Deserializer::ParseDepthwiseConvolution2d;
     m_ParserFunctions[Layer_DequantizeLayer]             = &Deserializer::ParseDequantize;
     m_ParserFunctions[Layer_DetectionPostProcessLayer]   = &Deserializer::ParseDetectionPostProcess;
@@ -258,6 +259,8 @@ Deserializer::LayerBaseRawPtr Deserializer::GetBaseLayer(const GraphPtr& graphPt
             return graphPtr->layers()->Get(layerIndex)->layer_as_ConstantLayer()->base();
         case Layer::Layer_Convolution2dLayer:
             return graphPtr->layers()->Get(layerIndex)->layer_as_Convolution2dLayer()->base();
+        case Layer::Layer_DepthToSpaceLayer:
+            return graphPtr->layers()->Get(layerIndex)->layer_as_DepthToSpaceLayer()->base();
         case Layer::Layer_DepthwiseConvolution2dLayer:
             return graphPtr->layers()->Get(layerIndex)->layer_as_DepthwiseConvolution2dLayer()->base();
         case Layer::Layer_DequantizeLayer:
@@ -1115,6 +1118,32 @@ void Deserializer::ParseConvolution2d(GraphPtr graph, unsigned int layerIndex)
     RegisterOutputSlots(graph, layerIndex, layer);
 }
 
+void Deserializer::ParseDepthToSpace(GraphPtr graph, unsigned int layerIndex)
+{
+    CHECK_LAYERS(graph, 0, layerIndex);
+
+    auto inputs = GetInputs(graph, layerIndex);
+    CHECK_VALID_SIZE(inputs.size(), 1);
+
+    auto outputs = GetOutputs(graph, layerIndex);
+    CHECK_VALID_SIZE(outputs.size(), 1);
+
+    auto fbDescriptor = graph->layers()->Get(layerIndex)->layer_as_DepthToSpaceLayer()->descriptor();
+
+    armnn::DepthToSpaceDescriptor descriptor;
+    descriptor.m_BlockSize  = fbDescriptor->blockSize();
+    descriptor.m_DataLayout = ToDataLayout(fbDescriptor->dataLayout());
+
+    auto layerName = GetLayerName(graph, layerIndex);
+    IConnectableLayer* layer = m_Network->AddDepthToSpaceLayer(descriptor, layerName.c_str());
+
+    armnn::TensorInfo outputInfo = ToTensorInfo(outputs[0]);
+    layer->GetOutputSlot(0).SetTensorInfo(outputInfo);
+
+    RegisterInputSlots(graph, layerIndex, layer);
+    RegisterOutputSlots(graph, layerIndex, layer);
+}
+
 void Deserializer::ParseDepthwiseConvolution2d(GraphPtr graph, unsigned int layerIndex)
 {
     CHECK_LAYERS(graph, 0, layerIndex);
diff --git a/src/armnnDeserializer/Deserializer.hpp b/src/armnnDeserializer/Deserializer.hpp
index 30af024f5e..0e2a34fc80 100644
--- a/src/armnnDeserializer/Deserializer.hpp
+++ b/src/armnnDeserializer/Deserializer.hpp
@@ -86,6 +86,7 @@ private:
     void ParseConcat(GraphPtr graph, unsigned int layerIndex);
     void ParseConstant(GraphPtr graph, unsigned int layerIndex);
     void ParseConvolution2d(GraphPtr graph, unsigned int layerIndex);
+    void ParseDepthToSpace(GraphPtr graph, unsigned int layerIndex);
     void ParseDepthwiseConvolution2d(GraphPtr graph, unsigned int layerIndex);
     void ParseDequantize(GraphPtr graph, unsigned int layerIndex);
     void ParseDetectionPostProcess(GraphPtr graph, unsigned int layerIndex);
diff --git a/src/armnnDeserializer/DeserializerSupport.md b/src/armnnDeserializer/DeserializerSupport.md
index 46bf92eadd..78ae51675f 100644
--- a/src/armnnDeserializer/DeserializerSupport.md
+++ b/src/armnnDeserializer/DeserializerSupport.md
@@ -15,6 +15,7 @@ The Arm NN SDK Deserialize parser currently supports the following layers:
 * Concat
 * Constant
 * Convolution2d
+* DepthToSpace
 * DepthwiseConvolution2d
 * Dequantize
 * DetectionPostProcess
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