!6146 Fix bug of quantize operator.
Merge pull request !6146 from wangshaocong/lite_quant
This commit is contained in:
commit
12f3665167
|
@ -39,7 +39,8 @@ constexpr uint32_t kDoubleNum = 2;
|
|||
constexpr uint32_t kMultiNum = 3;
|
||||
constexpr uint32_t kDimension_4d = 4;
|
||||
|
||||
const std::set<int> kSupportDataType = {kNumberTypeUInt8, kNumberTypeInt32, kNumberTypeFloat32, kNumberTypeFloat16};
|
||||
const std::set<int> kSupportDataType = {kNumberTypeUInt8, kNumberTypeInt8, kNumberTypeInt32,
|
||||
kNumberTypeFloat32, kNumberTypeFloat16};
|
||||
|
||||
#ifdef PRIMITIVE_WRITEABLE
|
||||
using TensorPtr = std::shared_ptr<mindspore::tensor::Tensor>;
|
||||
|
|
|
@ -147,6 +147,7 @@ kernel::LiteKernel *CpuCastFp32KernelCreator(const std::vector<lite::Tensor *> &
|
|||
|
||||
REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_Cast, CpuCastFp32KernelCreator)
|
||||
REG_KERNEL(kCPU, kNumberTypeUInt8, PrimitiveType_Cast, CpuCastFp32KernelCreator)
|
||||
REG_KERNEL(kCPU, kNumberTypeInt8, PrimitiveType_Cast, CpuCastFp32KernelCreator)
|
||||
REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_Cast, CpuCastFp32KernelCreator)
|
||||
#ifndef ENABLE_ARM64
|
||||
REG_KERNEL(kCPU, kNumberTypeFloat16, PrimitiveType_Cast, CpuCastFp32KernelCreator)
|
||||
|
|
|
@ -36,27 +36,37 @@ STATUS TfliteQuantizeParser::Parse(const std::unique_ptr<tflite::OperatorT> &tfl
|
|||
return RET_NULL_PTR;
|
||||
}
|
||||
|
||||
std::unique_ptr<schema::QuantDTypeCastT> attr = std::make_unique<schema::QuantDTypeCastT>();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "new op failed";
|
||||
return RET_NULL_PTR;
|
||||
}
|
||||
|
||||
const auto &in_tensor = tflite_tensors[tflite_op->inputs[0]];
|
||||
if (in_tensor == nullptr) {
|
||||
MS_LOG(ERROR) << "input tensor is null";
|
||||
return RET_NULL_PTR;
|
||||
}
|
||||
attr->srcT = GetTfliteDataType(in_tensor->type);
|
||||
const auto &out_tensor = tflite_tensors[tflite_op->outputs[0]];
|
||||
if (out_tensor == nullptr) {
|
||||
MS_LOG(ERROR) << "output tensor is null";
|
||||
return RET_NULL_PTR;
|
||||
}
|
||||
attr->dstT = GetTfliteDataType(out_tensor->type);
|
||||
|
||||
op->primitive->value.type = schema::PrimitiveType_QuantDTypeCast;
|
||||
op->primitive->value.value = attr.release();
|
||||
if (GetTfliteDataType(in_tensor->type) != kNumberTypeInt8) {
|
||||
std::unique_ptr<schema::QuantDTypeCastT> attr = std::make_unique<schema::QuantDTypeCastT>();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "new op failed";
|
||||
return RET_NULL_PTR;
|
||||
}
|
||||
attr->srcT = GetTfliteDataType(in_tensor->type);
|
||||
attr->dstT = GetTfliteDataType(out_tensor->type);
|
||||
op->primitive->value.type = schema::PrimitiveType_QuantDTypeCast;
|
||||
op->primitive->value.value = attr.release();
|
||||
} else {
|
||||
std::unique_ptr<schema::CastT> attr = std::make_unique<schema::CastT>();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "new op failed";
|
||||
return RET_NULL_PTR;
|
||||
}
|
||||
attr->srcT = GetTfliteDataType(in_tensor->type);
|
||||
attr->dstT = GetTfliteDataType(out_tensor->type);
|
||||
op->primitive->value.type = schema::PrimitiveType_Cast;
|
||||
op->primitive->value.value = attr.release();
|
||||
}
|
||||
|
||||
AddOpInput(op, tensors_id, tensors_format, tensors_id_map, tflite_op->inputs[0], tensors_id->size(),
|
||||
tflite_tensors.size(), schema::Format_NHWC);
|
||||
|
|
Loading…
Reference in New Issue