experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Check and report allocation failures in InitializeRuntimeTensor.

Browse Source 
PiperOrigin-RevId: 306497434
Change-Id: I1b94e40f5ea12e5f2c9fbb73be534b12016e76fe
This commit is contained in:
Robert David 2020-04-14 12:58:54 -07:00 committed by TensorFlower Gardener
parent 6f1bbf2cc2
commit e2cec54297
1 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
tensorflow/lite/micro/micro_allocator.cc
View File

@ -351,12 +351,29 @@ TfLiteStatus InitializeRuntimeTensor(
reinterpret_cast<TfLiteAffineQuantization*>(
allocator->AllocateFromTail(sizeof(TfLiteAffineQuantization),
alignof(TfLiteAffineQuantization)));
if (quantization == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate TfLiteAffineQuantization.\n");
return kTfLiteError;
}
quantization->zero_point =
reinterpret_cast<TfLiteIntArray*>(allocator->AllocateFromTail(
TfLiteIntArrayGetSizeInBytes(channels), alignof(TfLiteIntArray)));
if (quantization->zero_point == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate quantization->zero_point.\n");
return kTfLiteError;
}
quantization->scale = reinterpret_cast<TfLiteFloatArray*>(
allocator->AllocateFromTail(TfLiteFloatArrayGetSizeInBytes(channels),
alignof(TfLiteFloatArray)));
if (quantization->scale == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter,
"Unable to allocate quantization->scale.\n");
return kTfLiteError;
}
quantization->zero_point->size = channels;
quantization->scale->size = channels;
int* zero_point_data = quantization->zero_point->data;