IVGCVSW-2625 Support static quantization of softmax

Change-Id: I216344ee10bbffadb648d4aef1d9a0d9dbb4a341
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Signed-off-by: ruoyan01 <ruomei.yan@arm.com>

5 files changed, 75 insertions, 2 deletions

diff --git a/src/armnn/QuantizerVisitor.cpp b/src/armnn/QuantizerVisitor.cpp
index f5ff83c31f..08384e285d 100644
--- a/src/armnn/QuantizerVisitor.cpp
+++ b/src/armnn/QuantizerVisitor.cpp
@@ -177,4 +177,13 @@ void QuantizerVisitor::VisitConvolution2dLayer(const IConnectableLayer* layer,
     SetQuantizedInputConnections(layer, newLayer);
 }
 
-} // namespace armnn
+void QuantizerVisitor::VisitSoftmaxLayer(const IConnectableLayer* layer,
+                                         const SoftmaxDescriptor& softmaxDescriptor,
+                                         const char* name)
+{
+    IConnectableLayer* newLayer = m_QuantizedNetwork->AddSoftmaxLayer(softmaxDescriptor, name);
+    RecordLayer(layer, newLayer);
+    SetQuantizedInputConnections(layer, newLayer);
+}
+
+} //namespace armnn
diff --git a/src/armnn/QuantizerVisitor.hpp b/src/armnn/QuantizerVisitor.hpp
index 8d8b787d89..44ebc052b6 100644
--- a/src/armnn/QuantizerVisitor.hpp
+++ b/src/armnn/QuantizerVisitor.hpp
@@ -58,6 +58,10 @@ public:
                                  const Optional<ConstTensor>& biases,
                                  const char* name = nullptr) override;
 
+    void VisitSoftmaxLayer(const IConnectableLayer* layer,
+                           const SoftmaxDescriptor& softmaxDescriptor,
+                           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 6ae0d3314b..05117cd5cd 100644
--- a/src/armnn/StaticRangeVisitor.cpp
+++ b/src/armnn/StaticRangeVisitor.cpp
@@ -113,4 +113,12 @@ void StaticRangeVisitor::VisitFullyConnectedLayer(const IConnectableLayer *layer
     SetRange(layer, 0, -15.0f, 15.0f);
 }
 
+void StaticRangeVisitor::VisitSoftmaxLayer(const IConnectableLayer* layer,
+                                           const SoftmaxDescriptor& softmaxDescriptor,
+                                           const char* name)
+{
+    boost::ignore_unused(softmaxDescriptor);
+    SetRange(layer, 0, 0.f, 1.f);
+}
+
 } //namespace armnn
diff --git a/src/armnn/StaticRangeVisitor.hpp b/src/armnn/StaticRangeVisitor.hpp
index 887f3e9120..9c3a4f32c1 100644
--- a/src/armnn/StaticRangeVisitor.hpp
+++ b/src/armnn/StaticRangeVisitor.hpp
@@ -52,7 +52,9 @@ public:
                                   const ConstTensor& weights,
                                   const Optional<ConstTensor>& biases,
                                   const char *name) override;
-
+    void VisitSoftmaxLayer(const IConnectableLayer* layer,
+                           const SoftmaxDescriptor& softmaxDescriptor,
+                           const char* name = nullptr) override;
     /// Retrieve the default range
     MinMaxRange DefaultRange() const { return std::make_pair(-15.0f, 15.0f); }
 
diff --git a/src/armnn/test/QuantizerTest.cpp b/src/armnn/test/QuantizerTest.cpp
index 90935f37f8..c83d179961 100644
--- a/src/armnn/test/QuantizerTest.cpp
+++ b/src/armnn/test/QuantizerTest.cpp
@@ -629,5 +629,55 @@ BOOST_AUTO_TEST_CASE(QuantizeConvolution2dWithBiases)
     TestQuantizeConvolution2d(true);
 }
 
+class TestSoftmaxQuantization : public TestQuantization
+{
+public:
+    virtual void VisitSoftmaxLayer(const IConnectableLayer* layer,
+                                   const SoftmaxDescriptor& descriptor,
+                                   const char* name = nullptr)
+    {
+        TensorInfo info = layer->GetOutputSlot(0).GetTensorInfo();
+
+        BOOST_TEST((info.GetDataType() ==  DataType::QuantisedAsymm8));
+
+        BOOST_TEST((info.GetQuantizationOffset() == 0));
+
+        BOOST_CHECK_CLOSE(info.GetQuantizationScale(), 1.0f/255.0f, 0.000001f );
+    }
+};
+
+INetworkPtr CreateNetworkWithSoftmaxLayer(const SoftmaxDescriptor& descriptor)
+{
+    auto network = INetwork::Create();
+    // Add the layers
+    IConnectableLayer* input0 = network->AddInputLayer(0);
+    IConnectableLayer* softmax = network->AddSoftmaxLayer(descriptor);
+    IConnectableLayer* output = network->AddOutputLayer(2);
+
+    // Establish connections
+    input0->GetOutputSlot(0).Connect(softmax->GetInputSlot(0));
+    softmax->GetOutputSlot(0).Connect(output->GetInputSlot(0));
+
+    //Set TensorInfo
+    TensorShape shape{1U};
+    TensorInfo info(shape, DataType::Float32);
+    input0->GetOutputSlot(0).SetTensorInfo(info);
+    softmax->GetOutputSlot(0).SetTensorInfo(info);
+
+    return network;
+}
+
+BOOST_AUTO_TEST_CASE(QuantizeSoftmax)
+{
+    SoftmaxDescriptor descriptor;
+    descriptor.m_Beta = 1.0f;
+
+    auto network = CreateNetworkWithSoftmaxLayer(descriptor);
+
+    auto quantizedNetwork = INetworkQuantizer::Create(network.get())->ExportNetwork();
+    TestSoftmaxQuantization validator;
+    VisitLayersTopologically(quantizedNetwork.get(), validator);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace armnn