diff options
-rw-r--r-- | src/armnnSerializer/Serializer.cpp | 26 | ||||
-rw-r--r-- | src/armnnSerializer/Serializer.hpp | 13 |
2 files changed, 34 insertions, 5 deletions
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp index 3923f0f9f1..79ad7e19af 100644 --- a/src/armnnSerializer/Serializer.cpp +++ b/src/armnnSerializer/Serializer.cpp @@ -35,6 +35,23 @@ serializer::DataType GetFlatBufferDataType(DataType dataType) } } +uint32_t SerializerVisitor::GetSerializedId(unsigned int guid) +{ + std::pair<unsigned int, uint32_t> guidPair(guid, m_layerId); + + if (m_guidMap.empty()) + { + m_guidMap.insert(guidPair); + } + else if (m_guidMap.find(guid) == m_guidMap.end()) + { + guidPair.second = ++m_layerId; + m_guidMap.insert(guidPair); + return m_layerId; + } + return m_layerId; +} + // Build FlatBuffer for Input Layer void SerializerVisitor::VisitInputLayer(const IConnectableLayer* layer, LayerBindingId id, const char* name) { @@ -47,7 +64,7 @@ void SerializerVisitor::VisitInputLayer(const IConnectableLayer* layer, LayerBin id); // Push layer Guid to outputIds. - m_inputIds.push_back(layer->GetGuid()); + m_inputIds.push_back(GetSerializedId(layer->GetGuid())); // Create the FlatBuffer InputLayer auto flatBufferInputLayer = serializer::CreateInputLayer(m_flatBufferBuilder, flatBufferInputBindableBaseLayer); @@ -67,7 +84,7 @@ void SerializerVisitor::VisitOutputLayer(const IConnectableLayer* layer, LayerBi flatBufferOutputBaseLayer, id); // Push layer Guid to outputIds. - m_outputIds.push_back(layer->GetGuid()); + m_outputIds.push_back(GetSerializedId(layer->GetGuid())); // Create the FlatBuffer OutputLayer auto flatBufferOutputLayer = serializer::CreateOutputLayer(m_flatBufferBuilder, flatBufferOutputBindableBaseLayer); @@ -109,7 +126,7 @@ fb::Offset<serializer::LayerBase> SerializerVisitor::CreateLayerBase(const IConn std::vector<fb::Offset<serializer::OutputSlot>> outputSlots = CreateOutputSlots(layer); return serializer::CreateLayerBase(m_flatBufferBuilder, - layer->GetGuid(), + GetSerializedId(layer->GetGuid()), m_flatBufferBuilder.CreateString(layer->GetName()), layerType, m_flatBufferBuilder.CreateVector(inputSlots), @@ -137,7 +154,8 @@ std::vector<fb::Offset<serializer::InputSlot>> SerializerVisitor::CreateInputSlo const IOutputSlot* connection = inputSlot.GetConnection(); // Create FlatBuffer Connection - serializer::Connection conn(connection->GetOwningLayerGuid(), connection->CalculateIndexOnOwner()); + serializer::Connection conn(GetSerializedId(inputSlot.GetConnection()->GetOwningLayerGuid()), + connection->CalculateIndexOnOwner()); // Create FlatBuffer InputSlot inputSlots.push_back(serializer::CreateInputSlot(m_flatBufferBuilder, slotIndex, &conn)); } diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp index b799dbe380..8a509e880e 100644 --- a/src/armnnSerializer/Serializer.hpp +++ b/src/armnnSerializer/Serializer.hpp @@ -10,6 +10,8 @@ #include <armnnSerializer/ISerializer.hpp> #include <iostream> +#include <unordered_map> + #include <Schema_generated.h> namespace armnnSerializer @@ -18,7 +20,7 @@ namespace armnnSerializer class SerializerVisitor : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy> { public: - SerializerVisitor() {} + SerializerVisitor() : m_layerId(0) {}; ~SerializerVisitor() {} flatbuffers::FlatBufferBuilder& GetFlatBufferBuilder() @@ -65,6 +67,9 @@ private: /// Creates the serializer AnyLayer for the layer and adds it to m_serializedLayers. void CreateAnyLayer(const flatbuffers::Offset<void>& layer, const armnn::armnnSerializer::Layer serializerLayer); + ///Function which maps Guid to an index + uint32_t GetSerializedId(unsigned int guid); + /// Creates the serializer InputSlots for the layer. std::vector<flatbuffers::Offset<armnn::armnnSerializer::InputSlot>> CreateInputSlots( const armnn::IConnectableLayer* layer); @@ -84,6 +89,12 @@ private: /// Guids of all Output Layers required by the SerializedGraph. std::vector<unsigned int> m_outputIds; + + /// Mapped Guids of all Layers to match our index. + std::unordered_map<unsigned int, uint32_t > m_guidMap; + + /// layer within our FlatBuffer index. + uint32_t m_layerId; }; class Serializer : public ISerializer |