diff options
author | Teresa Charlin <teresa.charlinreyes@arm.com> | 2023-03-28 17:20:19 +0100 |
---|---|---|
committer | TeresaARM <teresa.charlinreyes@arm.com> | 2023-03-29 11:10:14 +0000 |
commit | 046e2cbc82ddac369f21bfa7c93375338c93afd7 (patch) | |
tree | f40e14a62db9cb8202e98b1c2dca31fe92072066 | |
parent | ad1b3d7518429e2d16a2695d9b0bbf81b6565ac9 (diff) | |
download | armnn-046e2cbc82ddac369f21bfa7c93375338c93afd7.tar.gz |
BugFix: Mean layer reading negative axis incorrectly in TFLite parser
Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I3ccfc2a4c95c57743927fb276faa756740f5c55a
-rw-r--r-- | src/armnnTfLiteParser/TfLiteParser.cpp | 36 |
1 files changed, 27 insertions, 9 deletions
diff --git a/src/armnnTfLiteParser/TfLiteParser.cpp b/src/armnnTfLiteParser/TfLiteParser.cpp index dc5afca30e..c787212359 100644 --- a/src/armnnTfLiteParser/TfLiteParser.cpp +++ b/src/armnnTfLiteParser/TfLiteParser.cpp @@ -2529,20 +2529,38 @@ void TfLiteParserImpl::ParseMean(size_t subgraphIndex, size_t operatorIndex) auto outputs = GetOutputs(m_Model, subgraphIndex, operatorIndex); CHECK_VALID_SIZE(outputs.size(), 1); - armnn::TensorInfo dimTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 1); - BufferRawPtr bufferPtr = GetBuffer(m_Model, inputs[1]->buffer); + TensorInfo inputTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 0); + TensorInfo dimTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 1); armnn::MeanDescriptor desc; - std::vector<unsigned int> axis(dimTensorInfo.GetNumElements()); - ::memcpy(axis.data(), bufferPtr->data.data(), dimTensorInfo.GetNumBytes()); - desc.m_Axis = axis; + BufferRawPtr axisBufferPtr = GetBuffer(m_Model, inputs[1]->buffer); + // Get const axis value from model and set it to descriptor. + if (axisBufferPtr != nullptr) + { + std::vector<int32_t> axisData(dimTensorInfo.GetNumElements()); + ::memcpy(axisData.data(), axisBufferPtr->data.data(), dimTensorInfo.GetNumBytes()); + + // Convert the axis to unsigned int and remove duplicates. + auto rank = static_cast<int32_t>(inputTensorInfo.GetNumDimensions()); + std::set<unsigned int> uniqueAxis; + std::transform(axisData.begin(), + axisData.end(), + std::inserter(uniqueAxis, uniqueAxis.begin()), + [rank](int i)->unsigned int{ + return static_cast<uint32_t>(((i + rank) % rank)); }); + desc.m_Axis.assign(uniqueAxis.begin(), uniqueAxis.end()); + } + else + { + for (uint32_t i = 0; i < inputTensorInfo.GetNumDimensions(); ++i) + { + desc.m_Axis.push_back(i); + } + } - armnn::TensorInfo inputTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 0); armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0], true); - desc.m_KeepDims = - inputTensorInfo.GetNumDimensions() == outputTensorInfo.GetNumDimensions() ? - true : false; + desc.m_KeepDims = inputTensorInfo.GetNumDimensions() == outputTensorInfo.GetNumDimensions() ? true : false; auto layerName = fmt::format("Mean:{}:{}", subgraphIndex, operatorIndex); IConnectableLayer* layer = m_Network->AddMeanLayer(desc, layerName.c_str()); |