wehub-resource-sync
89 lines
3.0 KiB
C++
89 lines
3.0 KiB
C++
//
|
|
// TFCustomDequantize.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/07/28.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../merge/MergeHelpers.hpp"
|
|
#include "MNN_generated.h"
|
|
#include "TFExtraManager.hpp"
|
|
#include <string>
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class CustomDequantizeTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
MNN_ASSERT(nullptr != op->main_as_Extra());
|
|
auto attr = op->main_as_Extra()->attr();
|
|
|
|
int nbit = 8;
|
|
float zero_point = 0.f;
|
|
float clamp_min = -128.0f;
|
|
float clamp_max = 127.0f;
|
|
int method;
|
|
std::vector<float> scale_val;
|
|
for (int i = 0; i < attr->size(); ++i) {
|
|
auto attr_value = attr->GetAs<Attribute>(i);
|
|
if (attr_value->key()->str() == "nbit") {
|
|
nbit = attr_value->i();
|
|
}
|
|
if (attr_value->key()->str() == "zero_point") {
|
|
zero_point = attr_value->f();
|
|
}
|
|
if (attr_value->key()->str() == "clamp_min") {
|
|
clamp_min = attr_value->f();
|
|
}
|
|
if (attr_value->key()->str() == "clamp_max") {
|
|
clamp_max = attr_value->f();
|
|
}
|
|
if (attr_value->key()->str() == "method") {
|
|
method = attr_value->i();
|
|
}
|
|
if (attr_value->key()->str() == "scale") {
|
|
auto* list_value = attr_value->list()->f();
|
|
scale_val.resize(list_value->size());
|
|
for (int i = 0; i < list_value->size(); ++i) {
|
|
scale_val[i] = list_value->Get(i);
|
|
}
|
|
}
|
|
}
|
|
|
|
VARP input = inputs[0];
|
|
std::unique_ptr<OpT> dequant_op(new OpT);
|
|
dequant_op->type = OpType_Int8ToFloat;
|
|
dequant_op->main.type = OpParameter_QuantizedFloatParam;
|
|
dequant_op->main.value = new QuantizedFloatParamT;
|
|
|
|
auto* dequant_param = dequant_op->main.AsQuantizedFloatParam();
|
|
dequant_param->nbits = nbit;
|
|
dequant_param->clampMin = clamp_min;
|
|
dequant_param->clampMax = clamp_max;
|
|
dequant_param->method = MNN::QuantizeAlgo(method);
|
|
dequant_param->tensorScale = {scale_val[0]};
|
|
EXPRP dequant_expr = Expr::create(dequant_op.get(), {input});
|
|
dequant_expr->setName(expr->name());
|
|
|
|
VARP dequant = Variable::create(dequant_expr);
|
|
dequant->setName(expr->name());
|
|
VARP output = helpers::ConvertLayout(dequant, NHWC, NC4HW4);
|
|
output->setName(expr->name() + "_convert_to_nhwc");
|
|
output->expr().first->setName(output->name());
|
|
return output->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
TFExtraManager::get()->insert("CustomDequantize",
|
|
std::shared_ptr<TFExtraManager::Transform>(new CustomDequantizeTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|