diff options
Diffstat (limited to 'src/armnnSerializer')
-rw-r--r-- | src/armnnSerializer/ArmnnSchema.fbs | 1 | ||||
-rw-r--r-- | src/armnnSerializer/ArmnnSchema_generated.h | 14 | ||||
-rw-r--r-- | src/armnnSerializer/Serializer.cpp | 4 |
3 files changed, 16 insertions, 3 deletions
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs index 40de3496b0..362dd5c653 100644 --- a/src/armnnSerializer/ArmnnSchema.fbs +++ b/src/armnnSerializer/ArmnnSchema.fbs @@ -495,6 +495,7 @@ table SoftmaxLayer { table SoftmaxDescriptor { beta:float; + axis:int = -1; } table DepthwiseConvolution2dLayer { diff --git a/src/armnnSerializer/ArmnnSchema_generated.h b/src/armnnSerializer/ArmnnSchema_generated.h index 7747f9edd9..b66bac69db 100644 --- a/src/armnnSerializer/ArmnnSchema_generated.h +++ b/src/armnnSerializer/ArmnnSchema_generated.h @@ -5126,14 +5126,19 @@ inline flatbuffers::Offset<SoftmaxLayer> CreateSoftmaxLayer( struct SoftmaxDescriptor FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SoftmaxDescriptorBuilder Builder; enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE { - VT_BETA = 4 + VT_BETA = 4, + VT_AXIS = 6 }; float beta() const { return GetField<float>(VT_BETA, 0.0f); } + int32_t axis() const { + return GetField<int32_t>(VT_AXIS, -1); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField<float>(verifier, VT_BETA) && + VerifyField<int32_t>(verifier, VT_AXIS) && verifier.EndTable(); } }; @@ -5145,6 +5150,9 @@ struct SoftmaxDescriptorBuilder { void add_beta(float beta) { fbb_.AddElement<float>(SoftmaxDescriptor::VT_BETA, beta, 0.0f); } + void add_axis(int32_t axis) { + fbb_.AddElement<int32_t>(SoftmaxDescriptor::VT_AXIS, axis, -1); + } explicit SoftmaxDescriptorBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -5159,8 +5167,10 @@ struct SoftmaxDescriptorBuilder { inline flatbuffers::Offset<SoftmaxDescriptor> CreateSoftmaxDescriptor( flatbuffers::FlatBufferBuilder &_fbb, - float beta = 0.0f) { + float beta = 0.0f, + int32_t axis = -1) { SoftmaxDescriptorBuilder builder_(_fbb); + builder_.add_axis(axis); builder_.add_beta(beta); return builder_.Finish(); } diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp index c08784352d..a9f7b7d307 100644 --- a/src/armnnSerializer/Serializer.cpp +++ b/src/armnnSerializer/Serializer.cpp @@ -1037,7 +1037,9 @@ void SerializerStrategy::SerializeSoftmaxLayer(const armnn::IConnectableLayer* l // Create the FlatBuffer SoftmaxDescriptor auto flatBufferSoftmaxDesc = - serializer::CreateSoftmaxDescriptor(m_flatBufferBuilder, softmaxDescriptor.m_Beta); + serializer::CreateSoftmaxDescriptor(m_flatBufferBuilder, + softmaxDescriptor.m_Beta, + softmaxDescriptor.m_Axis); // Create the FlatBuffer SoftmaxLayer auto flatBufferSoftmaxLayer = |