diff options
Diffstat (limited to 'delegate/src/Activation.hpp')
-rw-r--r-- | delegate/src/Activation.hpp | 104 |
1 files changed, 96 insertions, 8 deletions
diff --git a/delegate/src/Activation.hpp b/delegate/src/Activation.hpp index 1ffa7d7f8c..5e8d876110 100644 --- a/delegate/src/Activation.hpp +++ b/delegate/src/Activation.hpp @@ -5,7 +5,7 @@ #pragma once -#include <armnn/utility/IgnoreUnused.hpp> +#include "DelegateUtils.hpp" #include <tensorflow/lite/builtin_ops.h> #include <tensorflow/lite/c/builtin_op_data.h> @@ -15,20 +15,108 @@ namespace armnnDelegate { +TfLiteStatus ValidateActivationOperator(DelegateData& delegateData, + TfLiteContext* tfLiteContext, + const armnn::TensorInfo& inputInfo, + const armnn::TensorInfo& outputInfo, + armnn::ActivationDescriptor& activationDesc) +{ + bool isSupported = false; + auto validateFunc = [&](const armnn::TensorInfo& outInfo, bool& isSupported) + { + FORWARD_LAYER_SUPPORT_FUNC(__func__, + tfLiteContext, + IsActivationSupported, + delegateData.m_Backends, + isSupported, + inputInfo, + outputInfo, + activationDesc); + }; + + validateFunc(outputInfo, isSupported); + return isSupported ? kTfLiteOk : kTfLiteError; +} + TfLiteStatus VisitActivationOperator(DelegateData& delegateData, TfLiteContext* tfLiteContext, TfLiteNode* tfLiteNode, int nodeIndex, - int32_t comparisonOperatorCode) + int32_t operatorCode) { - armnn::IgnoreUnused(delegateData, - tfLiteContext, - tfLiteNode, - nodeIndex, - comparisonOperatorCode); + TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 1, nodeIndex)); + TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex)); + + const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors; + const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[0]]; + if (IsDynamicTensor(tfLiteInputTensor)) + { + TF_LITE_MAYBE_KERNEL_LOG( + tfLiteContext, + "TfLiteArmnnDelegate: Dynamic input tensors are not supported in node #%d: ", + nodeIndex); + return kTfLiteError; + } + const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]]; + if (IsDynamicTensor(tfLiteOutputTensor)) + { + TF_LITE_MAYBE_KERNEL_LOG( + tfLiteContext, + "TfLiteArmnnDelegate: Dynamic output tensors are not supported in node #%d: ", + nodeIndex); + return kTfLiteError; + } + + const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteInputTensor); + const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor); + + armnn::ActivationDescriptor activationDesc; + switch(operatorCode) + { + case kTfLiteBuiltinRelu: + { + activationDesc.m_Function = armnn::ActivationFunction::ReLu; + break; + } + case kTfLiteBuiltinRelu6: + { + activationDesc.m_Function = armnn::ActivationFunction::BoundedReLu; + activationDesc.m_A = 6.0f; + break; + } + case kTfLiteBuiltinLogistic: + { + activationDesc.m_Function = armnn::ActivationFunction::Sigmoid; + break; + } + case kTfLiteBuiltinTanh: + { + activationDesc.m_Function = armnn::ActivationFunction::TanH; + activationDesc.m_A = 1.0f; + activationDesc.m_B = 1.0f; + break; + } + default: + { + return kTfLiteError; + } + } + if (!delegateData.m_Network) + { + return ValidateActivationOperator(delegateData, + tfLiteContext, + inputTensorInfo, + outputTensorInfo, + activationDesc); + } + armnn::IConnectableLayer* activationLayer = delegateData.m_Network->AddActivationLayer(activationDesc); + ARMNN_ASSERT(activationLayer != nullptr); + armnn::IOutputSlot& outputSlot = activationLayer->GetOutputSlot(0); + outputSlot.SetTensorInfo(outputTensorInfo); - return kTfLiteError; + // Connect + return Connect(activationLayer, tfLiteNode, delegateData); } } // namespace armnnDelegate |