aboutsummaryrefslogtreecommitdiff
path: root/src/armnnSerializer/Serializer.cpp
diff options
context:
space:
mode:
Diffstat (limited to 'src/armnnSerializer/Serializer.cpp')
-rw-r--r--src/armnnSerializer/Serializer.cpp21
1 files changed, 21 insertions, 0 deletions
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index 79ad7e19af..ba4b36934c 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -119,6 +119,27 @@ void SerializerVisitor::VisitMultiplicationLayer(const IConnectableLayer* layer,
CreateAnyLayer(flatBufferMultiplicationLayer.o, serializer::Layer::Layer_MultiplicationLayer);
}
+// Build FlatBuffer for Softmax Layer
+void SerializerVisitor::VisitSoftmaxLayer(const IConnectableLayer* layer,
+ const SoftmaxDescriptor& softmaxDescriptor,
+ const char* name)
+{
+ // Create FlatBuffer BaseLayer
+ auto flatBufferSoftmaxBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_Softmax);
+
+ // Create the FlatBuffer SoftmaxDescriptor
+ auto flatBufferSoftmaxDesc =
+ serializer::CreateSoftmaxDescriptor(m_flatBufferBuilder, softmaxDescriptor.m_Beta);
+
+ // Create the FlatBuffer SoftmaxLayer
+ auto flatBufferSoftmaxLayer =
+ serializer::CreateSoftmaxLayer(m_flatBufferBuilder,
+ flatBufferSoftmaxBaseLayer,
+ flatBufferSoftmaxDesc);
+
+ CreateAnyLayer(flatBufferSoftmaxLayer.o, serializer::Layer::Layer_SoftmaxLayer);
+}
+
fb::Offset<serializer::LayerBase> SerializerVisitor::CreateLayerBase(const IConnectableLayer* layer,
const serializer::LayerType layerType)
{