src/armnn/layers/BatchNormalizationLayer.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. 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 <backendsCommon/CpuTensorHandle.hpp>
  11 #include <backendsCommon/WorkloadFactory.hpp>
  12
  13 namespace armnn
  14 {
  15
  16 BatchNormalizationLayer::BatchNormalizationLayer(const armnn::BatchNormalizationDescriptor& param, const char* name)
  17     : LayerWithParameters(1, 1, LayerType::BatchNormalization, param, name)
  18 {
  19 }
  20
  21 std::unique_ptr<IWorkload> BatchNormalizationLayer::CreateWorkload(const Graph& graph,
  22                                                                    const IWorkloadFactory& factory) const
  23 {
  24     // on this level constant data should not be released..
  25     BOOST_ASSERT_MSG(m_Mean != nullptr, "BatchNormalizationLayer: Mean data should not be null.");
  26     BOOST_ASSERT_MSG(m_Variance != nullptr, "BatchNormalizationLayer: Variance data should not be null.");
  27     BOOST_ASSERT_MSG(m_Beta != nullptr, "BatchNormalizationLayer: Beta data should not be null.");
  28     BOOST_ASSERT_MSG(m_Gamma != nullptr, "BatchNormalizationLayer: Gamma data should not be null.");
  29
  30     BatchNormalizationQueueDescriptor 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.CreateBatchNormalization(descriptor, PrepInfoAndDesc(descriptor, graph));
  38 }
  39
  40 BatchNormalizationLayer* BatchNormalizationLayer::Clone(Graph& graph) const
  41 {
  42     auto layer = CloneBase<BatchNormalizationLayer>(graph, m_Param, GetName());
  43
  44     layer->m_Mean = m_Mean ? std::make_unique<ScopedCpuTensorHandle>(*m_Mean) : nullptr;
  45     layer->m_Variance = m_Variance ? std::make_unique<ScopedCpuTensorHandle>(*m_Variance) : nullptr;
  46     layer->m_Beta = m_Beta ? std::make_unique<ScopedCpuTensorHandle>(*m_Beta) : nullptr;
  47     layer->m_Gamma = m_Gamma ? std::make_unique<ScopedCpuTensorHandle>(*m_Gamma) : nullptr;
  48
  49     return std::move(layer);
  50 }
  51
  52 void BatchNormalizationLayer::ValidateTensorShapesFromInputs()
  53 {
  54     VerifyLayerConnections(1, CHECK_LOCATION());
  55
  56     auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
  57
  58     BOOST_ASSERT(inferredShapes.size() == 1);
  59
  60     ConditionalThrowIfNotEqual<LayerValidationException>(
  61         "BatchNormalizationLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
  62         GetOutputSlot(0).GetTensorInfo().GetShape(),
  63         inferredShapes[0]);
  64
  65 }
  66
  67 Layer::ConstantTensors BatchNormalizationLayer::GetConstantTensorsByRef()
  68 {
  69     return {m_Mean, m_Variance, m_Beta, m_Gamma};
  70 }
  71
  72 } // namespace armnn