wehub-resource-sync
70 lines
2.6 KiB
C++
70 lines
2.6 KiB
C++
//
|
|
// SoftmaxTflite.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "liteOpConverter.hpp"
|
|
|
|
DECLARE_OP_COVERTER(SoftmaxTflite);
|
|
|
|
MNN::OpType SoftmaxTflite::opType(int quantizedModel) {
|
|
if (quantizedModel)
|
|
return MNN::OpType_QuantizedSoftmax;
|
|
return MNN::OpType_Softmax;
|
|
}
|
|
MNN::OpParameter SoftmaxTflite::type(int quantizedModel) {
|
|
if (quantizedModel)
|
|
return MNN::OpParameter_QuantizedSoftmax;
|
|
return MNN::OpParameter_Axis;
|
|
}
|
|
|
|
void SoftmaxTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>& tfliteOp,
|
|
const std::vector<std::unique_ptr<tflite::TensorT>>& tfliteTensors,
|
|
const std::vector<std::unique_ptr<tflite::BufferT>>& tfliteModelBuffer,
|
|
const std::vector<std::unique_ptr<tflite::OperatorCodeT>>& tfliteOpSet, int quantizedModel) {
|
|
DCHECK(tfliteOp->inputs.size() == 1) << "Tflite Softmax input ERROR!";
|
|
const auto& tfliteSoftmaxOption = tfliteOp->builtin_options.AsSoftmaxOptions();
|
|
const int outputIndex = tfliteOp->outputs[0];
|
|
const auto& outputTensor = tfliteTensors[outputIndex];
|
|
if (outputTensor->type == tflite::TensorType_INT8) {
|
|
quantizedModel = 2;
|
|
dstOp->type = MNN::OpType_Softmax;
|
|
dstOp->main.type = MNN::OpParameter_Axis;
|
|
} else if (outputTensor->type == tflite::TensorType_UINT8) {
|
|
quantizedModel = 1;
|
|
dstOp->type = MNN::OpType_QuantizedSoftmax;
|
|
dstOp->main.type = MNN::OpParameter_QuantizedSoftmax;
|
|
} else {
|
|
MNN_ASSERT(outputTensor->type == tflite::TensorType_FLOAT32);
|
|
quantizedModel = 0;
|
|
dstOp->type = MNN::OpType_Softmax;
|
|
dstOp->main.type = MNN::OpParameter_Axis;
|
|
}
|
|
|
|
if (quantizedModel == 1) {
|
|
auto softmaxParamQuan = new MNN::QuantizedSoftmaxT();
|
|
softmaxParamQuan->beta = tfliteSoftmaxOption->beta;
|
|
// input
|
|
const int inputIndex = tfliteOp->inputs[0];
|
|
const auto& inputTensor = tfliteTensors[inputIndex];
|
|
softmaxParamQuan->inputScale = inputTensor->quantization->scale[0];
|
|
dstOp->main.value = softmaxParamQuan;
|
|
} else {
|
|
auto paramFloat = new MNN::AxisT;
|
|
paramFloat->axis = -1;
|
|
dstOp->main.value = paramFloat;
|
|
}
|
|
// set input output index
|
|
dstOp->inputIndexes.resize(1);
|
|
dstOp->outputIndexes.resize(1);
|
|
dstOp->inputIndexes[0] = tfliteOp->inputs[0];
|
|
dstOp->outputIndexes[0] = tfliteOp->outputs[0];
|
|
}
|
|
|
|
using namespace tflite;
|
|
REGISTER_CONVERTER(SoftmaxTflite, BuiltinOperator_SOFTMAX);
|