src/armnn/layers/MeanLayer.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5
   6 #include "MeanLayer.hpp"
   7 #include "LayerCloneBase.hpp"
   8
   9 #include <backendsCommon/CpuTensorHandle.hpp>
  10 #include <backendsCommon/WorkloadData.hpp>
  11 #include <backendsCommon/WorkloadFactory.hpp>
  12
  13 #include <cstring>
  14
  15 namespace armnn
  16 {
  17
  18 MeanLayer::MeanLayer(const armnn::MeanDescriptor& param, const char* name)
  19     : LayerWithParameters(1, 1, LayerType::Mean, param, name)
  20 {}
  21
  22 std::unique_ptr<IWorkload> MeanLayer::CreateWorkload(const armnn::Graph& graph,
  23                                                      const armnn::IWorkloadFactory& factory) const
  24 {
  25     MeanQueueDescriptor descriptor;
  26     descriptor.m_Parameters.m_Axis = m_Param.m_Axis;
  27     descriptor.m_Parameters.m_KeepDims = m_Param.m_KeepDims;
  28
  29     return factory.CreateMean(descriptor, PrepInfoAndDesc(descriptor, graph));
  30 }
  31
  32 MeanLayer* MeanLayer::Clone(Graph& graph) const
  33 {
  34     auto layer = CloneBase<MeanLayer>(graph, m_Param, GetName());
  35
  36     layer->m_Param.m_Axis = m_Param.m_Axis;
  37     layer->m_Param.m_KeepDims = m_Param.m_KeepDims;
  38
  39     return std::move(layer);
  40 }
  41
  42 void MeanLayer::ValidateTensorShapesFromInputs()
  43 {
  44     VerifyLayerConnections(1, CHECK_LOCATION());
  45
  46     const TensorInfo& input = GetInputSlot(0).GetConnection()->GetTensorInfo();
  47
  48     BOOST_ASSERT_MSG(input.GetNumDimensions() > 0 && input.GetNumDimensions() <= MaxNumOfTensorDimensions,
  49                      "MeanLayer: Mean supports up to 4D input.");
  50
  51     unsigned int rank = input.GetNumDimensions();
  52     unsigned int outputRank = 0;
  53
  54     // Calculate output dimension
  55     if (m_Param.m_KeepDims)
  56     {
  57         outputRank = rank;
  58     }
  59     else if (m_Param.m_Axis.empty())
  60     {
  61         outputRank = 1;
  62     }
  63     else if (m_Param.m_Axis.size() >= input.GetNumDimensions())
  64     {
  65         throw LayerValidationException("MeanLayer: Dimensions to reduce can not be bigger than input dimensions");
  66     }
  67     else
  68     {
  69         outputRank = input.GetNumDimensions() - boost::numeric_cast<unsigned int>(m_Param.m_Axis.size());
  70         if (outputRank == 0)
  71         {
  72             outputRank = 1;
  73         }
  74     }
  75
  76     std::vector<unsigned int> dimSizes(outputRank, 1);
  77     if (!m_Param.m_Axis.empty())
  78     {
  79         // Skip the dimension that has been reduced unless keepDims is true.
  80         unsigned int outputIndex = 0;
  81         for (unsigned int i = 0; i < input.GetNumDimensions(); ++i)
  82         {
  83             if (std::find(m_Param.m_Axis.begin(), m_Param.m_Axis.end(), i) == m_Param.m_Axis.end())
  84             {
  85                 dimSizes[outputIndex] = boost::numeric_cast<unsigned int>(input.GetShape()[i]);
  86                 ++outputIndex;
  87             }
  88             else if (m_Param.m_KeepDims)
  89             {
  90                 dimSizes[outputIndex] = 1;
  91                 ++outputIndex;
  92             }
  93         }
  94     }
  95     const TensorShape& inferredShape = TensorShape(outputRank, dimSizes.data());
  96
  97     ConditionalThrowIfNotEqual<LayerValidationException>(
  98         "MeanLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
  99         GetOutputSlot(0).GetTensorInfo().GetShape(),
 100         inferredShape);
 101 }
 102
 103 } // namespace armnn