IVGCVSW-2613 Support static quantization of BatchToSpace

Change-Id: I44b12c5c246b7aacc789420dbe55a16efaab6f98

5 files changed, 58 insertions, 0 deletions

diff --git a/src/armnn/QuantizerVisitor.cpp b/src/armnn/QuantizerVisitor.cpp
index c0cf0dfc5b..437d7b90d6 100644
--- a/src/armnn/QuantizerVisitor.cpp
+++ b/src/armnn/QuantizerVisitor.cpp
@@ -298,4 +298,13 @@ void QuantizerVisitor::VisitStridedSliceLayer(const IConnectableLayer* layer,
     SetQuantizedInputConnections(layer, newLayer);
 }
 
+void QuantizerVisitor::VisitBatchToSpaceNdLayer(const IConnectableLayer* layer,
+                                                const BatchToSpaceNdDescriptor& batchToSpaceNdDescriptor,
+                                                const char* name)
+{
+    IConnectableLayer* newLayer = m_QuantizedNetwork->AddBatchToSpaceNdLayer(batchToSpaceNdDescriptor, name);
+    RecordLayer(layer, newLayer);
+    SetQuantizedInputConnections(layer, newLayer);
+}
+
 } //namespace armnn
diff --git a/src/armnn/QuantizerVisitor.hpp b/src/armnn/QuantizerVisitor.hpp
index 6954ede8c7..2682663047 100644
--- a/src/armnn/QuantizerVisitor.hpp
+++ b/src/armnn/QuantizerVisitor.hpp
@@ -104,6 +104,10 @@ public:
                                 const StridedSliceDescriptor& stridedSliceDescriptor,
                                 const char* name = nullptr) override;
 
+    void VisitBatchToSpaceNdLayer(const IConnectableLayer* layer,
+                                  const BatchToSpaceNdDescriptor& batchToSpaceNdDescriptor,
+                                  const char* name = nullptr) override;
+
     /// Extract the quantized network
     INetworkPtr RetrieveFinalNetwork() { return std::move(m_QuantizedNetwork); }
 
diff --git a/src/armnn/StaticRangeVisitor.cpp b/src/armnn/StaticRangeVisitor.cpp
index 766b4ad3db..815730be31 100644
--- a/src/armnn/StaticRangeVisitor.cpp
+++ b/src/armnn/StaticRangeVisitor.cpp
@@ -234,4 +234,13 @@ void StaticRangeVisitor::VisitStridedSliceLayer(const IConnectableLayer* layer,
     ForwardParentParameters(layer);
 }
 
+void StaticRangeVisitor::VisitBatchToSpaceNdLayer(const IConnectableLayer* layer,
+                                                  const BatchToSpaceNdDescriptor& batchToSpaceNdDescriptor,
+                                                  const char* name)
+{
+    boost::ignore_unused(batchToSpaceNdDescriptor);
+    boost::ignore_unused(name);
+    ForwardParentParameters(layer);
+}
+
 } //namespace armnn
diff --git a/src/armnn/StaticRangeVisitor.hpp b/src/armnn/StaticRangeVisitor.hpp
index 1002c860a5..5a16e184d6 100644
--- a/src/armnn/StaticRangeVisitor.hpp
+++ b/src/armnn/StaticRangeVisitor.hpp
@@ -95,6 +95,10 @@ public:
                                 const StridedSliceDescriptor& stridedSliceDescriptor,
                                 const char* name = nullptr) override;
 
+    void VisitBatchToSpaceNdLayer(const IConnectableLayer* layer,
+                                  const BatchToSpaceNdDescriptor& batchToSpaceNdDescriptor,
+                                  const char* name = nullptr) override;
+
 private:
     /// Set the range for an output slot on a layer
     void SetRange(const IConnectableLayer* layer, unsigned int outputIdx, float min, float max);
diff --git a/src/armnn/test/QuantizerTest.cpp b/src/armnn/test/QuantizerTest.cpp
index 3edea9cc6e..548203a6a9 100644
--- a/src/armnn/test/QuantizerTest.cpp
+++ b/src/armnn/test/QuantizerTest.cpp
@@ -1176,6 +1176,7 @@ BOOST_AUTO_TEST_CASE(QuantizeStridedSlice)
     // Add the layer under test
     StridedSliceDescriptor stridedSliceDesc;
     IConnectableLayer* stridedSlice = network->AddStridedSliceLayer(stridedSliceDesc);
+
     CompleteLeakyReluNetwork(network.get(), activation, stridedSlice, info);
 
     auto quantizedNetwork = INetworkQuantizer::Create(network.get())->ExportNetwork();
@@ -1183,5 +1184,36 @@ BOOST_AUTO_TEST_CASE(QuantizeStridedSlice)
     VisitLayersTopologically(quantizedNetwork.get(), validator);
 }
 
+BOOST_AUTO_TEST_CASE(QuantizeBatchToSpace)
+{
+    class TestBatchToSpaceQuantization : public TestLeakyReLuActivationQuantization
+    {
+    public:
+        void VisitBatchToSpaceNdLayer(const IConnectableLayer* layer,
+                                      const BatchToSpaceNdDescriptor& batchToSpaceNdDescriptor,
+                                      const char* name = nullptr) override
+        {
+            CheckForwardedQuantizationSettings(layer);
+        }
+    };
+
+    INetworkPtr network = INetwork::Create();
+
+    TensorShape shape{1U};
+    TensorInfo info(shape, DataType::Float32);
+
+    IConnectableLayer* activation = CreateStartOfLeakyReluNetwork(network.get(), info);
+
+    // Add the layer under test
+    BatchToSpaceNdDescriptor descriptor;
+    IConnectableLayer* batchToSpace = network->AddBatchToSpaceNdLayer(descriptor);
+
+    CompleteLeakyReluNetwork(network.get(), activation, batchToSpace, info);
+
+    auto quantizedNetwork = INetworkQuantizer::Create(network.get())->ExportNetwork();
+    TestBatchToSpaceQuantization validator;
+    VisitLayersTopologically(quantizedNetwork.get(), validator);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace armnn