xref: /aosp_15_r20/external/armnn/src/armnn/layers/BatchNormalizationLayer.cpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
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