1 //
2 // Copyright © 2017 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #include "SpaceToDepthLayer.hpp"
7 #include "LayerCloneBase.hpp"
8
9 #include <armnn/TypesUtils.hpp>
10 #include <armnn/utility/IgnoreUnused.hpp>
11 #include <armnnUtils/DataLayoutIndexed.hpp>
12
13 #include <backendsCommon/WorkloadData.hpp>
14 #include <backendsCommon/WorkloadFactory.hpp>
15
16 #include <numeric>
17
18 using namespace armnnUtils;
19
20 namespace armnn
21 {
22
SpaceToDepthLayer(const SpaceToDepthDescriptor param,const char * name)23 SpaceToDepthLayer::SpaceToDepthLayer(const SpaceToDepthDescriptor param, const char* name)
24 : LayerWithParameters(1, 1, LayerType::SpaceToDepth, param, name)
25 {}
26
CreateWorkload(const IWorkloadFactory & factory) const27 std::unique_ptr<IWorkload> SpaceToDepthLayer::CreateWorkload(const IWorkloadFactory& factory) const
28 {
29 SpaceToDepthQueueDescriptor descriptor;
30 descriptor.m_Parameters.m_BlockSize = m_Param.m_BlockSize;
31 descriptor.m_Parameters.m_DataLayout = m_Param.m_DataLayout;
32
33 SetAdditionalInfo(descriptor);
34
35 return factory.CreateSpaceToDepth(descriptor, PrepInfoAndDesc(descriptor));
36 }
37
Clone(Graph & graph) const38 SpaceToDepthLayer* SpaceToDepthLayer::Clone(Graph& graph) const
39 {
40 IgnoreUnused(graph);
41 return CloneBase<SpaceToDepthLayer>(graph, m_Param, GetName());
42 }
43
InferOutputShapes(const std::vector<TensorShape> & inputShapes) const44 std::vector<TensorShape> SpaceToDepthLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
45 {
46 ARMNN_ASSERT(inputShapes.size() == 1);
47
48 TensorShape inputShape = inputShapes[0];
49 TensorShape outputShape(inputShape);
50
51 DataLayoutIndexed dimensionIndices{m_Param.m_DataLayout};
52 unsigned int hIndex = dimensionIndices.GetHeightIndex();
53 unsigned int wIndex = dimensionIndices.GetWidthIndex();
54 unsigned int cIndex = dimensionIndices.GetChannelsIndex();
55
56 outputShape[hIndex] = inputShape[hIndex] / m_Param.m_BlockSize;
57 outputShape[wIndex] = inputShape[wIndex] / m_Param.m_BlockSize;
58
59 outputShape[cIndex] = inputShape[cIndex] * m_Param.m_BlockSize * m_Param.m_BlockSize;
60
61 return std::vector<TensorShape>({ outputShape });
62 }
63
ValidateTensorShapesFromInputs()64 void SpaceToDepthLayer::ValidateTensorShapesFromInputs()
65 {
66 VerifyLayerConnections(1, CHECK_LOCATION());
67
68 const TensorShape& outputShape = GetOutputSlot(0).GetTensorInfo().GetShape();
69
70 VerifyShapeInferenceType(outputShape, m_ShapeInferenceMethod);
71
72 std::vector<TensorShape> inferredShapes = InferOutputShapes({
73 GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
74
75 ARMNN_ASSERT(inferredShapes.size() == 1);
76
77 ValidateAndCopyShape(outputShape, inferredShapes[0], m_ShapeInferenceMethod, "SpaceToDepthLayer");
78 }
79
Accept(ILayerVisitor & visitor) const80 void SpaceToDepthLayer::Accept(ILayerVisitor& visitor) const
81 {
82 visitor.VisitSpaceToDepthLayer(this, GetParameters(), GetName());
83 }
84
85 } // namespace armnn
86