diff options
Diffstat (limited to 'src/armnnSerializer')
-rw-r--r-- | src/armnnSerializer/ArmnnSchema.fbs | 15 | ||||
-rw-r--r-- | src/armnnSerializer/Serializer.cpp | 29 | ||||
-rw-r--r-- | src/armnnSerializer/Serializer.hpp | 4 | ||||
-rw-r--r-- | src/armnnSerializer/SerializerSupport.md | 1 | ||||
-rw-r--r-- | src/armnnSerializer/test/SerializerTests.cpp | 46 |
5 files changed, 94 insertions, 1 deletions
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs index c1e4e3826a..7d64dca6f3 100644 --- a/src/armnnSerializer/ArmnnSchema.fbs +++ b/src/armnnSerializer/ArmnnSchema.fbs @@ -94,7 +94,8 @@ enum LayerType : uint { Permute = 10, FullyConnected = 11, Constant = 12, - SpaceToBatchNd = 13 + SpaceToBatchNd = 13, + BatchToSpaceNd = 14 } // Base layer table to be used as part of other layers @@ -266,9 +267,21 @@ table SpaceToBatchNdDescriptor { dataLayout:DataLayout; } +table BatchToSpaceNdLayer { + base:LayerBase; + descriptor:BatchToSpaceNdDescriptor; +} + +table BatchToSpaceNdDescriptor { + blockShape:[uint]; + crops:[uint]; + dataLayout:DataLayout; +} + union Layer { ActivationLayer, AdditionLayer, + BatchToSpaceNdLayer, ConstantLayer, Convolution2dLayer, DepthwiseConvolution2dLayer, diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp index 1c1cefd455..dc25bab0c6 100644 --- a/src/armnnSerializer/Serializer.cpp +++ b/src/armnnSerializer/Serializer.cpp @@ -141,6 +141,35 @@ void SerializerVisitor::VisitAdditionLayer(const armnn::IConnectableLayer* layer CreateAnyLayer(flatBufferAdditionLayer.o, serializer::Layer::Layer_AdditionLayer); } +// Build FlatBuffer for BatchToSpaceNd Layer +void SerializerVisitor::VisitBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer, + const armnn::BatchToSpaceNdDescriptor& descriptor, + const char* name) +{ + // Create FlatBuffer BaseLayer + auto flatBufferBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_BatchToSpaceNd); + + std::vector<unsigned int> crops; + crops.reserve(descriptor.m_Crops.size() * 2); + for (auto& crop : descriptor.m_Crops) + { + crops.push_back(crop.first); + crops.push_back(crop.second); + } + + auto flatBufferDescriptor = + CreateBatchToSpaceNdDescriptor(m_flatBufferBuilder, + m_flatBufferBuilder.CreateVector(descriptor.m_BlockShape), + m_flatBufferBuilder.CreateVector(crops), + GetFlatBufferDataLayout(descriptor.m_DataLayout)); + + auto flatBufferLayer = serializer::CreateBatchToSpaceNdLayer(m_flatBufferBuilder, + flatBufferBaseLayer, + flatBufferDescriptor); + + CreateAnyLayer(flatBufferLayer.o, serializer::Layer::Layer_BatchToSpaceNdLayer); +} + // Build FlatBuffer for Constant Layer void SerializerVisitor::VisitConstantLayer(const armnn::IConnectableLayer* layer, const armnn::ConstTensor& input, diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp index a86e04f05d..f8e2b16bf9 100644 --- a/src/armnnSerializer/Serializer.hpp +++ b/src/armnnSerializer/Serializer.hpp @@ -49,6 +49,10 @@ public: void VisitAdditionLayer(const armnn::IConnectableLayer* layer, const char* name = nullptr) override; + void VisitBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer, + const armnn::BatchToSpaceNdDescriptor& descriptor, + const char* name = nullptr) override; + void VisitConstantLayer(const armnn::IConnectableLayer* layer, const armnn::ConstTensor& input, const char* = nullptr) override; diff --git a/src/armnnSerializer/SerializerSupport.md b/src/armnnSerializer/SerializerSupport.md index 776f6129bd..6340dd2e9a 100644 --- a/src/armnnSerializer/SerializerSupport.md +++ b/src/armnnSerializer/SerializerSupport.md @@ -8,6 +8,7 @@ The Arm NN SDK Serializer currently supports the following layers: * Activation * Addition +* BatchToSpaceNd * Constant * Convolution2d * DepthwiseConvolution2d diff --git a/src/armnnSerializer/test/SerializerTests.cpp b/src/armnnSerializer/test/SerializerTests.cpp index f64c14a184..8ce90499df 100644 --- a/src/armnnSerializer/test/SerializerTests.cpp +++ b/src/armnnSerializer/test/SerializerTests.cpp @@ -682,4 +682,50 @@ BOOST_AUTO_TEST_CASE(SerializeDeserializeSpaceToBatchNd) outputTensorInfo.GetShape()); } +BOOST_AUTO_TEST_CASE(SerializeDeserializeBatchToSpaceNd) +{ + class VerifyBatchToSpaceNdName : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy> + { + public: + void VisitBatchToSpaceNdLayer(const armnn::IConnectableLayer*, + const armnn::BatchToSpaceNdDescriptor& descriptor, + const char* name) override + { + BOOST_TEST(name == "BatchToSpaceNdLayer"); + } + }; + + unsigned int inputShape[] = {4, 1, 2, 2}; + unsigned int outputShape[] = {1, 1, 4, 4}; + + armnn::BatchToSpaceNdDescriptor desc; + desc.m_DataLayout = armnn::DataLayout::NCHW; + desc.m_BlockShape = {2, 2}; + desc.m_Crops = {{0, 0}, {0, 0}}; + + auto inputTensorInfo = armnn::TensorInfo(4, inputShape, armnn::DataType::Float32); + auto outputTensorInfo = armnn::TensorInfo(4, outputShape, armnn::DataType::Float32); + + armnn::INetworkPtr network = armnn::INetwork::Create(); + armnn::IConnectableLayer* const inputLayer = network->AddInputLayer(0); + armnn::IConnectableLayer* const batchToSpaceNdLayer = network->AddBatchToSpaceNdLayer(desc, "BatchToSpaceNdLayer"); + armnn::IConnectableLayer* const outputLayer = network->AddOutputLayer(0); + + inputLayer->GetOutputSlot(0).Connect(batchToSpaceNdLayer->GetInputSlot(0)); + inputLayer->GetOutputSlot(0).SetTensorInfo(inputTensorInfo); + batchToSpaceNdLayer->GetOutputSlot(0).Connect(outputLayer->GetInputSlot(0)); + batchToSpaceNdLayer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo); + + armnn::INetworkPtr deserializedNetwork = DeserializeNetwork(SerializeNetwork(*network)); + BOOST_CHECK(deserializedNetwork); + + VerifyBatchToSpaceNdName nameChecker; + deserializedNetwork->Accept(nameChecker); + + CheckDeserializedNetworkAgainstOriginal(*network, + *deserializedNetwork, + inputTensorInfo.GetShape(), + outputTensorInfo.GetShape()); +} + BOOST_AUTO_TEST_SUITE_END() |