diff options
-rw-r--r-- | src/armnnDeserializer/Deserializer.cpp | 4 | ||||
-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 |
4 files changed, 19 insertions, 4 deletions
diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp index bfd4f6b560..8b1e9b970c 100644 --- a/src/armnnDeserializer/Deserializer.cpp +++ b/src/armnnDeserializer/Deserializer.cpp @@ -1652,7 +1652,8 @@ void IDeserializer::DeserializerImpl::ParseFill(GraphPtr graph, unsigned int lay CHECK_VALID_SIZE(outputs.size(), 1); auto layerName = GetLayerName(graph, layerIndex); - armnn::FillDescriptor descriptor(1.0f); + armnn::FillDescriptor descriptor; + descriptor.m_Value = graph->layers()->Get(layerIndex)->layer_as_FillLayer()->descriptor()->value(); IConnectableLayer* layer = m_Network->AddFillLayer(descriptor, layerName.c_str()); armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0]); @@ -2484,6 +2485,7 @@ void IDeserializer::DeserializerImpl::ParseSoftmax(GraphPtr graph, unsigned int armnn::SoftmaxDescriptor descriptor; descriptor.m_Beta = graph->layers()->Get(layerIndex)->layer_as_SoftmaxLayer()->descriptor()->beta(); + descriptor.m_Axis = graph->layers()->Get(layerIndex)->layer_as_SoftmaxLayer()->descriptor()->axis(); auto layerName = GetLayerName(graph, layerIndex); IConnectableLayer* layer = m_Network->AddSoftmaxLayer(descriptor, layerName.c_str()); 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 = |