IVGCVSW-2615 Support static quantization of Pooling2d

Change-Id: Ica9c48adad741840b201047bb65541a19aac17f7
Signed-off-by: FrancisMurtagh <francis.murtagh@arm.com>

5 files changed, 77 insertions, 0 deletions

diff --git a/src/armnn/QuantizerVisitor.cpp b/src/armnn/QuantizerVisitor.cpp
index e808d4799d..850a3c1fae 100644
--- a/src/armnn/QuantizerVisitor.cpp
+++ b/src/armnn/QuantizerVisitor.cpp
@@ -224,6 +224,15 @@ void QuantizerVisitor::VisitSpaceToBatchNdLayer(const IConnectableLayer* layer,
     SetQuantizedInputConnections(layer, newLayer);
 }
 
+void QuantizerVisitor::VisitPooling2dLayer(const IConnectableLayer* layer,
+                                           const Pooling2dDescriptor& pooling2dDescriptor,
+                                           const char* name)
+{
+    IConnectableLayer* newLayer = m_QuantizedNetwork->AddPooling2dLayer(pooling2dDescriptor, name);
+    RecordLayer(layer, newLayer);
+    SetQuantizedInputConnections(layer, newLayer);
+}
+
 void QuantizerVisitor::VisitSoftmaxLayer(const IConnectableLayer* layer,
                                          const SoftmaxDescriptor& softmaxDescriptor,
                                          const char* name)
diff --git a/src/armnn/QuantizerVisitor.hpp b/src/armnn/QuantizerVisitor.hpp
index a39e6f13ad..8829bdfaa8 100644
--- a/src/armnn/QuantizerVisitor.hpp
+++ b/src/armnn/QuantizerVisitor.hpp
@@ -76,6 +76,10 @@ public:
                                   const SpaceToBatchNdDescriptor& spaceToBatchNdDescriptor,
                                   const char* name = nullptr) override;
 
+    void VisitPooling2dLayer(const IConnectableLayer* layer,
+                             const Pooling2dDescriptor& pooling2dDescriptor,
+                             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 6fc2bc6762..b1cbb2d574 100644
--- a/src/armnn/StaticRangeVisitor.cpp
+++ b/src/armnn/StaticRangeVisitor.cpp
@@ -134,6 +134,14 @@ void StaticRangeVisitor::VisitSpaceToBatchNdLayer(const IConnectableLayer* layer
     ForwardParentParameters(layer);
 }
 
+void StaticRangeVisitor::VisitPooling2dLayer(const IConnectableLayer* layer,
+                                             const Pooling2dDescriptor& pooling2dDescriptor,
+                                             const char* name)
+{
+    boost::ignore_unused(pooling2dDescriptor);
+    ForwardParentParameters(layer);
+}
+
 void StaticRangeVisitor::VisitSoftmaxLayer(const IConnectableLayer* layer,
                                            const SoftmaxDescriptor& softmaxDescriptor,
                                            const char* name)
diff --git a/src/armnn/StaticRangeVisitor.hpp b/src/armnn/StaticRangeVisitor.hpp
index ef3f4e4073..a7028d7a51 100644
--- a/src/armnn/StaticRangeVisitor.hpp
+++ b/src/armnn/StaticRangeVisitor.hpp
@@ -63,6 +63,10 @@ public:
                                   const SpaceToBatchNdDescriptor& spaceToBatchNdDescriptor,
                                   const char* name = nullptr) override;
 
+    void VisitPooling2dLayer(const IConnectableLayer* layer,
+                             const Pooling2dDescriptor& pooling2dDescriptor,
+                             const char* name) override;
+
     void VisitSoftmaxLayer(const IConnectableLayer* layer,
                            const SoftmaxDescriptor& softmaxDescriptor,
                            const char* name = nullptr) override;
diff --git a/src/armnn/test/QuantizerTest.cpp b/src/armnn/test/QuantizerTest.cpp
index 657a87abbc..5fcfc5d12e 100644
--- a/src/armnn/test/QuantizerTest.cpp
+++ b/src/armnn/test/QuantizerTest.cpp
@@ -879,5 +879,57 @@ BOOST_AUTO_TEST_CASE(QuantizeSpaceToBatch)
     VisitLayersTopologically(quantizedNetwork.get(), validator);
 }
 
+class TestPooling2dQuantization : public TestLeakyReLuActivationQuantization
+{
+public:
+    virtual void VisitPooling2dLayer(const IConnectableLayer* layer,
+                                     const Pooling2dDescriptor& desc,
+                                     const char* name = nullptr)
+    {
+        TensorInfo info = layer->GetOutputSlot(0).GetTensorInfo();
+
+        BOOST_TEST((info.GetDataType() == DataType::QuantisedAsymm8));
+
+        BOOST_TEST((info.GetQuantizationOffset() == 64));
+
+        // Based off parent LeakyReLu [-5.f, 15.f]
+        BOOST_CHECK_CLOSE(info.GetQuantizationScale(), 20.0f/255.0f, 0.000001f);
+    }
+};
+
+BOOST_AUTO_TEST_CASE(QuantizePooling2d)
+{
+    auto network = INetwork::Create();
+
+    TensorShape shape{1U};
+    TensorInfo info(shape, DataType::Float32);
+
+    Pooling2dDescriptor desc;
+    ActivationDescriptor activationDescriptor;
+    activationDescriptor.m_Function = ActivationFunction::LeakyReLu;
+    activationDescriptor.m_A        = 3.5f;
+    activationDescriptor.m_B        = -10.0f;
+
+    // Add the layers
+    IConnectableLayer* input0 = network->AddInputLayer(0);
+    IConnectableLayer* activation = network->AddActivationLayer(activationDescriptor);
+    IConnectableLayer* pooling2d = network->AddPooling2dLayer(desc);
+    IConnectableLayer* output = network->AddOutputLayer(3);
+
+    // Establish connections
+    input0->GetOutputSlot(0).Connect(activation->GetInputSlot(0));
+    activation->GetOutputSlot(0).Connect(pooling2d->GetInputSlot(0));
+    pooling2d->GetOutputSlot(0).Connect(output->GetInputSlot(0));
+
+    //Set TensorInfo
+    input0->GetOutputSlot(0).SetTensorInfo(info);
+    activation->GetOutputSlot(0).SetTensorInfo(info);
+    pooling2d->GetOutputSlot(0).SetTensorInfo(info);
+
+    auto quantizedNetwork = INetworkQuantizer::Create(network.get())->ExportNetwork();
+    TestPooling2dQuantization validator;
+    VisitLayersTopologically(quantizedNetwork.get(), validator);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace armnn