1 //
2 // Copyright © 2017 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 #include "BatchNormalizationLayer.hpp"
6
7 #include "LayerCloneBase.hpp"
8
9 #include <armnn/TypesUtils.hpp>
10 #include <armnn/backends/TensorHandle.hpp>
11 #include <armnn/backends/WorkloadFactory.hpp>
12
13 namespace armnn
14 {
15
BatchNormalizationLayer(const armnn::BatchNormalizationDescriptor & param,const char * name)16 BatchNormalizationLayer::BatchNormalizationLayer(const armnn::BatchNormalizationDescriptor& param, const char* name)
17 : LayerWithParameters(1, 1, LayerType::BatchNormalization, param, name)
18 {
19 }
20
CreateWorkload(const IWorkloadFactory & factory) const21 std::unique_ptr<IWorkload> BatchNormalizationLayer::CreateWorkload(const IWorkloadFactory& factory) const
22 {
23 // on this level constant data should not be released..
24 ARMNN_ASSERT_MSG(m_Mean != nullptr, "BatchNormalizationLayer: Mean data should not be null.");
25 ARMNN_ASSERT_MSG(m_Variance != nullptr, "BatchNormalizationLayer: Variance data should not be null.");
26 ARMNN_ASSERT_MSG(m_Beta != nullptr, "BatchNormalizationLayer: Beta data should not be null.");
27 ARMNN_ASSERT_MSG(m_Gamma != nullptr, "BatchNormalizationLayer: Gamma data should not be null.");
28
29 BatchNormalizationQueueDescriptor descriptor;
30 SetAdditionalInfo(descriptor);
31
32 descriptor.m_Mean = m_Mean.get();
33 descriptor.m_Variance = m_Variance.get();
34 descriptor.m_Beta = m_Beta.get();
35 descriptor.m_Gamma = m_Gamma.get();
36
37 return factory.CreateWorkload(LayerType::BatchNormalization, descriptor, PrepInfoAndDesc(descriptor));
38 }
39
Clone(Graph & graph) const40 BatchNormalizationLayer* BatchNormalizationLayer::Clone(Graph& graph) const
41 {
42 auto layer = CloneBase<BatchNormalizationLayer>(graph, m_Param, GetName());
43
44 layer->m_Mean = m_Mean ? m_Mean : nullptr;
45 layer->m_Variance = m_Variance ? m_Variance : nullptr;
46 layer->m_Beta = m_Beta ? m_Beta : nullptr;
47 layer->m_Gamma = m_Gamma ? m_Gamma : nullptr;
48
49 return std::move(layer);
50 }
51
ValidateTensorShapesFromInputs()52 void BatchNormalizationLayer::ValidateTensorShapesFromInputs()
53 {
54 VerifyLayerConnections(1, CHECK_LOCATION());
55
56 const TensorShape& outputShape = GetOutputSlot(0).GetTensorInfo().GetShape();
57
58 VerifyShapeInferenceType(outputShape, m_ShapeInferenceMethod);
59
60 auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
61
62 ARMNN_ASSERT(inferredShapes.size() == 1);
63
64 ValidateAndCopyShape(outputShape, inferredShapes[0], m_ShapeInferenceMethod, "BatchNormalizationLayer");
65
66 }
67
GetConstantTensorsByRef() const68 Layer::ImmutableConstantTensors BatchNormalizationLayer::GetConstantTensorsByRef() const
69 {
70 // For API stability DO NOT ALTER order and add new members to the end of vector
71 return {m_Mean, m_Variance, m_Beta, m_Gamma};
72 }
73
ExecuteStrategy(IStrategy & strategy) const74 void BatchNormalizationLayer::ExecuteStrategy(IStrategy& strategy) const
75 {
76 ManagedConstTensorHandle managedMean(m_Mean);
77 ManagedConstTensorHandle managedVariance(m_Variance);
78 ManagedConstTensorHandle managedBeta(m_Beta);
79 ManagedConstTensorHandle managedGamma(m_Gamma);
80
81 std::vector<armnn::ConstTensor> constTensors { { managedMean.GetTensorInfo(), managedMean.Map() },
82 { managedVariance.GetTensorInfo(), managedVariance.Map() },
83 { managedBeta.GetTensorInfo(), managedBeta.Map() },
84 { managedGamma.GetTensorInfo(), managedGamma.Map() } };
85
86 strategy.ExecuteStrategy(this, GetParameters(), constTensors, GetName());
87 }
88
89 } // namespace armnn
90