在AllocateTensors()函数中添加条件判断,以确保输入tensor指针不为空。示例代码如下所示:
void Interpreter::AllocateTensors() {
TensorAllocator* allocator = GetAllocator(0);
for (int i = 0; i < tensors_.size(); ++i) {
TfLiteTensor* t = &tensors_[i];
if (t->allocation_type != kNoAllocation) {
if (t->data.raw == nullptr || !allocator) {
// 若数据指针为空或分配器未定义,则给出错误提示信息。
LOG(ERROR) << "AllocateTensors: nullptr tensor or allocator";
return;
}
if (t->bytes == 0 || TfLiteIntArrayProduct(t->dims) == 0) {
// 零尺寸的tensor视为空tensor,不进行内存分配。
t->data.raw = nullptr;
continue;
}
size_t alignment = 8;
t->bytes = (t->bytes + alignment - 1) & ~(alignment - 1);
if (!allocator->Allocate(t)) {
// 内存分配失败,给出错误提示信息。
LOG(ERROR) << "AllocateTensors: failed to allocate tensor";
return;
}
}
}
in_flight_execution_plan_->UpdateInputOutputTensorPointers(tensors_);
}
这样,就会在输入tensor为NULL时避免空指针解引用的问题。