Skip to content

[Converter:Bugfix] Fix MNNConvert crash on malformed tflite models (#4796) #4823

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fshuang8299 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[Converter:Bugfix] Fix MNNConvert crash on malformed tflite models (#4796) #4823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion source/shape/ShapeConcat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@ class ConcatSizeComputer : public SizeComputer {
virtual bool onComputeSize(const MNN::Op* op, const std::vector<Tensor*>& inputs,
const std::vector<Tensor*>& outputs) const override {
MNN_ASSERT(1 == outputs.size());
// MNN_ASSERT(inputs.size() >= 2);
if (inputs.empty()) {
MNN_ERROR("Concat op has no input\n");
return false;
}
auto& ob = outputs[0]->buffer();
int basicAxis = 0;
if (op->type() == OpType_Concat) {
Expand All @@ -30,6 +33,10 @@ class ConcatSizeComputer : public SizeComputer {
// Concat-inputs may have scalar which should be delete
for (const auto& input : inputs) {
auto inputDimensions = input->buffer().dimensions;
if (inputDimensions <= 0 || inputDimensions > MNN_MAX_TENSOR_DIM) {
MNN_ERROR("Concat op input has invalid rank %d\n", inputDimensions);
return false;
}

// Tensor might be zeros size, but some dims may not be zero. should concat as usual.

Expand All @@ -39,12 +46,20 @@ class ConcatSizeComputer : public SizeComputer {
if (axis < 0) {
axis = inputDimensions + axis;
}
if (axis < 0 || axis >= inputDimensions) {
MNN_ERROR("Concat op axis %d out of range for %d dims\n", axis, inputDimensions);
return false;
}
break;
}


int sum = 0;
for (auto t : inputs) {
if (axis >= t->dimensions()) {
MNN_ERROR("Concat op axis %d out of range for %d dims\n", axis, t->dimensions());
return false;
}
sum += t->buffer().dim[axis].extent;
ob.type = t->buffer().type;
for (int i = 0; i < t->dimensions(); ++i) {
Expand Down
104 changes: 84 additions & 20 deletions tools/converter/source/tflite/ConvolutionTflite.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
//

#include <stdio.h>
#include <limits>

#include "TfliteUtils.hpp"
#include "liteOpConverter.hpp"
Expand Down Expand Up @@ -35,6 +36,11 @@ void Conv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>
const int inputSize = tfliteOp->inputs.size();
DCHECK(inputSize == 2 || inputSize == 3) << "tflite Conv2D input ERROR! ";
const auto& tfliteConvOption = tfliteOp->builtin_options.AsConv2DOptions();
if (nullptr == tfliteConvOption) {
DLOG(ERROR) << "CONV_2D operator carries no Conv2DOptions";
dstOp->type = MNN::OpType_MAX;
return;
}
const int inputIndex = tfliteOp->inputs[0];
const int weightIndex = tfliteOp->inputs[1];
const int outputIndex = tfliteOp->outputs[0];
Expand All @@ -60,12 +66,27 @@ void Conv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>
int group = 1;
// co kh kw ci
const auto& weightShape = weightTensor->shape;
DCHECK(weightShape.size() == 4) << "Conv2D weight ERROR!";
if (4 != weightShape.size()) {
DLOG(ERROR) << "CONV_2D weight shape is not 4-D";
dstOp->type = MNN::OpType_MAX;
return;
}
const int co = weightShape[0];
const int kh = weightShape[1];
const int kw = weightShape[2];
const int ci = weightShape[3];
const int weightSize = co * kh * kw * ci;
if (co <= 0 || kh <= 0 || kw <= 0 || ci <= 0) {
DLOG(ERROR) << "CONV_2D weight shape contains non-positive dimension";
dstOp->type = MNN::OpType_MAX;
return;
}
const int64_t weightSize64 = (int64_t)co * kh * kw * ci;
if (weightSize64 <= 0 || weightSize64 > std::numeric_limits<int>::max()) {
DLOG(ERROR) << "CONV_2D weight size overflow: " << co << "x" << kh << "x" << kw << "x" << ci;
dstOp->type = MNN::OpType_MAX;
return;
}
const int weightSize = (int)weightSize64;
if (inputShape.size() == 4 && inputShape[3] > ci) {
group = inputShape[3] / ci;
}
Expand Down Expand Up @@ -160,11 +181,14 @@ void Conv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>
conv2dParamQuan->biasQuantizedParam = std::unique_ptr<MNN::QuantizedParamT>(new MNN::QuantizedParamT);
conv2dParamQuan->biasQuantizedParam->zeroPoint = biasTensor->quantization->zero_point[0];
conv2dParamQuan->biasQuantizedParam->scale = biasTensor->quantization->scale[0];
DCHECK(biasData.size() / 4 == co) << "Bias Data ERROR";
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
std::vector<int32_t> biasInt32Vec(realBiasDataPtr, realBiasDataPtr + co);
conv2dParamQuan->bias = biasInt32Vec;
if (biasData.size() >= sizeof(int32_t) * co) {
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
std::vector<int32_t> biasInt32Vec(realBiasDataPtr, realBiasDataPtr + co);
conv2dParamQuan->bias = biasInt32Vec;
} else {
DLOG(ERROR) << "CONV_2D bias buffer is too small, ignore bias";
}
}

conv2dParamQuan->activationType = (MNN::FusedActivation)tfliteConvOption->fused_activation_function;
Expand Down Expand Up @@ -255,10 +279,15 @@ void Conv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>
convolution2DQuant->bias.resize(co);
if (inputSize == 3) {
const auto& biasTensor = tfliteTensors[tfliteOp->inputs[2]];
auto bias = reinterpret_cast<const int*>(tfliteModelBuffer[biasTensor->buffer]->data.data());
// int to float
for (int i = 0; i < co; i++) {
convolution2DQuant->bias[i] = bias[i] * (scaleIn * alpha[i]);
const auto& biasRaw = tfliteModelBuffer[biasTensor->buffer]->data;
auto bias = reinterpret_cast<const int*>(biasRaw.data());
if (biasRaw.size() >= sizeof(int) * co) {
// int to float
for (int i = 0; i < co; i++) {
convolution2DQuant->bias[i] = bias[i] * (scaleIn * alpha[i]);
}
} else {
DLOG(ERROR) << "CONV_2D bias buffer is too small, ignore bias";
}
}
dstOp->main.value = convolution2DQuant.release();
Expand Down Expand Up @@ -323,8 +352,12 @@ void Conv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT>
std::vector<float> biasData(co, 0.0f);
if (inputSize == 3) {
const auto& biasTensor = tfliteTensors[tfliteOp->inputs[2]];
auto biasDataPtr = reinterpret_cast<const float*>(tfliteModelBuffer[biasTensor->buffer]->data.data());
::memcpy(biasData.data(), biasDataPtr, sizeof(float) * co);
const auto& biasRaw = tfliteModelBuffer[biasTensor->buffer]->data;
if (biasRaw.data() != nullptr && biasRaw.size() >= sizeof(float) * co) {
::memcpy(biasData.data(), biasRaw.data(), sizeof(float) * co);
} else {
DLOG(ERROR) << "CONV_2D bias buffer is too small, ignore bias";
}
}
convolution2DFloat->bias = biasData;
dstOp->main.value = convolution2DFloat.release();
Expand Down Expand Up @@ -365,32 +398,58 @@ void TransposeConvTflite::run(MNN::OpT *dstOp, const std::unique_ptr<tflite::Ope
}
*/
const auto& tfliteConvOption = tfliteOp->builtin_options.AsTransposeConvOptions();
if (nullptr == tfliteConvOption) {
DLOG(ERROR) << "TRANSPOSE_CONV operator carries no TransposeConvOptions";
dstOp->type = MNN::OpType_MAX;
return;
}
// weight index
const int weightIndex = tfliteOp->inputs[1];
const auto& weightTensor = tfliteTensors[weightIndex];
// co kh kw ci
const auto& weightShape = weightTensor->shape;
DCHECK(weightShape.size() == 4) << "Conv2D weight ERROR!";
if (4 != weightShape.size()) {
DLOG(ERROR) << "TRANSPOSE_CONV weight shape is not 4-D";
dstOp->type = MNN::OpType_MAX;
return;
}
const int co = weightShape[0];
const int kh = weightShape[1];
const int kw = weightShape[2];
const int ci = weightShape[3];
const int weightSize = co * kh * kw * ci;
if (co <= 0 || kh <= 0 || kw <= 0 || ci <= 0) {
DLOG(ERROR) << "TRANSPOSE_CONV weight shape contains non-positive dimension";
dstOp->type = MNN::OpType_MAX;
return;
}
const int64_t weightSize64 = (int64_t)co * kh * kw * ci;
if (weightSize64 <= 0 || weightSize64 > std::numeric_limits<int>::max()) {
DLOG(ERROR) << "TRANSPOSE_CONV weight size overflow: " << co << "x" << kh << "x" << kw << "x" << ci;
dstOp->type = MNN::OpType_MAX;
return;
}
const int weightSize = (int)weightSize64;
{
auto convolution2DFloat = new MNN::Convolution2DT;
// weight
std::vector<float> weightData;
weightData.resize(weightSize);
auto originalWeightPtr = reinterpret_cast<const float*>(tfliteModelBuffer[weightTensor->buffer]->data.data());
convertDataFormatTflite(originalWeightPtr, weightData.data(), kh, kw, ci, co, true);
if (!convertDataFormatTflite(originalWeightPtr, weightData.data(), kh, kw, ci, co, true)) {
DLOG(ERROR) << "TRANSPOSE_CONV weight data is invalid";
dstOp->type = MNN::OpType_MAX;
return;
}
convolution2DFloat->weight = weightData;
// bias
std::vector<float> biasData(co, 0.0f);
if (inputSize == 4) {
const auto& biasTensor = tfliteTensors[tfliteOp->inputs[2]];
auto biasDataPtr = reinterpret_cast<const float*>(tfliteModelBuffer[biasTensor->buffer]->data.data());
if(biasDataPtr){
::memcpy(biasData.data(), biasDataPtr, sizeof(float) * co);
const auto& biasRaw = tfliteModelBuffer[biasTensor->buffer]->data;
if (biasRaw.data() != nullptr && biasRaw.size() >= sizeof(float) * co) {
::memcpy(biasData.data(), biasRaw.data(), sizeof(float) * co);
} else {
DLOG(ERROR) << "TRANSPOSE_CONV bias buffer is too small, ignore bias";
}
}
convolution2DFloat->bias = biasData;
Expand Down Expand Up @@ -440,11 +499,16 @@ void FullConnectedTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::Ope
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) {
const auto& option = tfliteOp->builtin_options.AsFullyConnectedOptions();
if (nullptr == option) {
DLOG(ERROR) << "FULLY_CONNECTED operator carries no FullyConnectedOptions";
dstOp->type = MNN::OpType_MAX;
return;
}
dstOp->main.value = new MNN::ExtraT;
auto dstP = dstOp->main.AsExtra();
dstP->engine = "Tflite";
dstP->type = "FULL_CONNECT";
const auto& option = tfliteOp->builtin_options.AsFullyConnectedOptions();
dstP->attr.resize(3);
dstP->attr[0].reset(new MNN::AttributeT);
dstP->attr[0]->key = "keep_num_dims";
Expand Down
85 changes: 60 additions & 25 deletions tools/converter/source/tflite/DepthwiseConv2DTflite.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
//

#include <stdio.h>
#include <limits>

#include "TfliteUtils.hpp"
#include "liteOpConverter.hpp"
Expand Down Expand Up @@ -78,13 +79,33 @@ void DepthwiseConv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::O
const auto& weightTensor = tfliteTensors[weightIndex];
// co kh kw ci
const auto& weightShape = weightTensor->shape;
DCHECK(weightShape.size() == 4) << "Conv2D weight ERROR!";
if (4 != weightShape.size()) {
DLOG(ERROR) << "DEPTHWISE_CONV_2D weight shape is not 4-D";
dstOp->type = MNN::OpType_MAX;
return;
}
// const int co = weightShape[0];
const int kh = weightShape[1];
const int kw = weightShape[2];
const int ci = weightShape[3];
const int weightSize = kh * kw * ci;
if (kh <= 0 || kw <= 0 || ci <= 0) {
DLOG(ERROR) << "DEPTHWISE_CONV_2D weight shape contains non-positive dimension";
dstOp->type = MNN::OpType_MAX;
return;
}
const int64_t weightSize64 = (int64_t)kh * kw * ci;
if (weightSize64 <= 0 || weightSize64 > std::numeric_limits<int>::max()) {
DLOG(ERROR) << "DEPTHWISE_CONV_2D weight size overflow: " << kh << "x" << kw << "x" << ci;
dstOp->type = MNN::OpType_MAX;
return;
}
const int weightSize = (int)weightSize64;
const auto& tfliteConvOption = tfliteOp->builtin_options.AsDepthwiseConv2DOptions();
if (nullptr == tfliteConvOption) {
DLOG(ERROR) << "DEPTHWISE_CONV_2D operator carries no DepthwiseConv2DOptions";
dstOp->type = MNN::OpType_MAX;
return;
}
if (weightTensor->type == tflite::TensorType_INT8) {
quantizedModel = 2;
dstOp->type = MNN::OpType_ConvolutionDepthwise;
Expand Down Expand Up @@ -115,30 +136,39 @@ void DepthwiseConv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::O
::memset(depthwiseConv2dParamFloat->bias.data(), 0, outputCount * sizeof(float));
if (inputSize == 3) {
const auto& biasTensor = tfliteTensors[tfliteOp->inputs[2]];
if (biasTensor->quantization->scale.size() == 1) {
auto scale = biasTensor->quantization->scale[0];
auto zero = biasTensor->quantization->zero_point[0];;
const auto& biasData = tfliteModelBuffer[biasTensor->buffer]->data;
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
for (int i=0; i<outputCount; ++i) {
depthwiseConv2dParamFloat->bias[i] = (float)(realBiasDataPtr[i] - zero) * scale;
const auto& biasData = tfliteModelBuffer[biasTensor->buffer]->data;
if (biasData.size() >= sizeof(int32_t) * outputCount) {
if (biasTensor->quantization->scale.size() == 1) {
auto scale = biasTensor->quantization->scale[0];
auto zero = biasTensor->quantization->zero_point[0];
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
for (int i = 0; i < outputCount; ++i) {
depthwiseConv2dParamFloat->bias[i] = (float)(realBiasDataPtr[i] - zero) * scale;
}
} else {
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
for (int i = 0; i < outputCount; ++i) {
depthwiseConv2dParamFloat->bias[i] =
(float)(realBiasDataPtr[i] - biasTensor->quantization->zero_point[i]) *
biasTensor->quantization->scale[i];
}
}
} else {
const auto& biasData = tfliteModelBuffer[biasTensor->buffer]->data;
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
for (int i=0; i<outputCount; ++i) {
depthwiseConv2dParamFloat->bias[i] = (float)(realBiasDataPtr[i] - biasTensor->quantization->zero_point[i]) * biasTensor->quantization->scale[i];
}
DLOG(ERROR) << "DEPTHWISE_CONV_2D bias buffer is too small, ignore bias";
}
}
// Weight
// Transpose first
std::vector<int8_t> transposeWeight(kw * kh * ci);
const auto& weightData = tfliteModelBuffer[weightTensor->buffer]->data;
auto weightDataPtr = (int8_t*)weightData.data();

if (weightDataPtr == nullptr || weightData.size() < (size_t)kw * kh * ci) {
DLOG(ERROR) << "DEPTHWISE_CONV_2D INT8 weight buffer is too small";
dstOp->type = MNN::OpType_MAX;
return;
}
for (int i=0; i<ci; ++i) {
for (int j=0; j<kw*kh; ++j) {
transposeWeight[i*kw*kh+j] = weightDataPtr[i+j*ci];
Expand Down Expand Up @@ -193,11 +223,14 @@ void DepthwiseConv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::O

auto shape = biasTensor->shape;

DCHECK(biasData.size() / 4 == ci) << "Bias Data ERROR";
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
std::vector<int32_t> biasInt32Vec(realBiasDataPtr, realBiasDataPtr + ci);
depthwiseConv2dParamQuan->bias = biasInt32Vec;
if (biasData.size() >= sizeof(int32_t) * ci) {
auto biasDataPtr = biasData.data();
const int32_t* realBiasDataPtr = (int32_t*)biasDataPtr;
std::vector<int32_t> biasInt32Vec(realBiasDataPtr, realBiasDataPtr + ci);
depthwiseConv2dParamQuan->bias = biasInt32Vec;
} else {
DLOG(ERROR) << "DEPTHWISE_CONV_2D bias buffer is too small, ignore bias";
}
}
depthwiseConv2dParamQuan->activationType =
static_cast<MNN::FusedActivation>(tfliteConvOption->fused_activation_function);
Expand All @@ -216,11 +249,13 @@ void DepthwiseConv2DTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::O
// bias
if (inputSize == 3) {
const auto& biasTensor = tfliteTensors[tfliteOp->inputs[2]];
auto originalBiasPtr = reinterpret_cast<const float*>(tfliteModelBuffer[biasTensor->buffer]->data.data());
if (originalBiasPtr) {
const auto& biasRaw = tfliteModelBuffer[biasTensor->buffer]->data;
if (biasRaw.data() != nullptr && biasRaw.size() >= sizeof(float) * ci) {
std::vector<float> biasData(ci, 0.0f);
::memcpy(biasData.data(), originalBiasPtr, sizeof(float) * ci);
::memcpy(biasData.data(), biasRaw.data(), sizeof(float) * ci);
depthwiseConv2dParamFloat->bias = biasData;
} else {
DLOG(ERROR) << "DEPTHWISE_CONV_2D bias buffer is too small, ignore bias";
}
}
depthwiseConv2dParamFloat->common = std::move(dstCommon);
Expand Down
5 changes: 5 additions & 0 deletions tools/converter/source/tflite/PoolingTflite.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,11 @@ void PoolingTflite::run(MNN::OpT* dstOp, const std::unique_ptr<tflite::OperatorT
const std::vector<std::unique_ptr<tflite::BufferT>>& tfliteModelBuffer,
const std::vector<std::unique_ptr<tflite::OperatorCodeT>>& tfliteOpSet, int quantizedModel) {
const auto& tflitePoolOption = tfliteOp->builtin_options.AsPool2DOptions();
if (nullptr == tflitePoolOption) {
DLOG(ERROR) << "AVERAGE_POOL_2D/MAX_POOL_2D operator carries no Pool2DOptions";
dstOp->type = MNN::OpType_MAX;
return;
}
const int outputIndex = tfliteOp->outputs[0];
const auto& outputTensor = tfliteTensors[outputIndex];
if (outputTensor->type == tflite::TensorType_INT8) {
Expand Down
Loading