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

Cleanup and refactor all allocations in MicroAllocator to function calls.

Browse Source 
This change is a precursor to adding a new memory logging MicroAllocator subclass that will enable TFLM to keep track of tensor arena tail allocations. Outside of moving all arena allocations to utility methods - I also cleaned up the organization of the methods inside of the cc file.

PiperOrigin-RevId: 313242666
Change-Id: Icddcc07187419fe314bc57708170cda8cd35690a
This commit is contained in:
Nick Kreeger 2020-05-26 12:12:38 -07:00 committed by TensorFlower Gardener
parent e76bff1201
commit 8c8dc2699b
4 changed files with 202 additions and 148 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -258,7 +258,7 @@ TfLiteStatus CommitPlan(ErrorReporter* error_reporter, MemoryPlanner* planner,
namespace internal {
TfLiteStatus InitializeRuntimeTensor(
TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
SimpleMemoryAllocator* allocator, const tflite::Tensor& flatbuffer_tensor,
const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers,
ErrorReporter* error_reporter, TfLiteTensor* result) {
@ -380,58 +380,9 @@ TfLiteStatus InitializeRuntimeTensor(
}
return kTfLiteOk;
}
} // namespace internal
TfLiteStatus MicroAllocator::Init() {
auto* subgraphs = model_->subgraphs();
if (subgraphs->size() != 1) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Only 1 subgraph is currently supported.\n");
return kTfLiteError;
}
subgraph_ = (*subgraphs)[0];
context_->tensors_size = subgraph_->tensors()->size();
context_->tensors =
reinterpret_cast<TfLiteTensor*>(memory_allocator_->AllocateFromTail(
sizeof(TfLiteTensor) * context_->tensors_size,
alignof(TfLiteTensor)));
if (context_->tensors == nullptr) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Failed to allocate memory for context->tensors, %d bytes required",
sizeof(TfLiteTensor) * context_->tensors_size);
return kTfLiteError;
}
// Initialize runtime tensors in context_ using the flatbuffer.
for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) {
TfLiteStatus status = internal::InitializeRuntimeTensor(
memory_allocator_, *subgraph_->tensors()->Get(i), model_->buffers(),
error_reporter_, &context_->tensors[i]);
if (status != kTfLiteOk) {
TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d",
i);
return kTfLiteError;
}
}
return kTfLiteOk;
}
size_t MicroAllocator::used_bytes() const {
if (active_) {
return 0;
}
TF_LITE_REPORT_ERROR(error_reporter_, "Total buffer usage: %d bytes",
memory_allocator_->GetUsedBytes());
TF_LITE_REPORT_ERROR(error_reporter_, "Head usage: %d bytes",
memory_allocator_->GetHeadUsedBytes());
TF_LITE_REPORT_ERROR(error_reporter_, "Tail usage: %d bytes",
memory_allocator_->GetTailUsedBytes());
return memory_allocator_->GetUsedBytes();
}
MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
uint8_t* tensor_arena, size_t arena_size,
ErrorReporter* error_reporter)
@ -450,7 +401,8 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
// destructed as it's the root allocator.
memory_allocator_ = CreateInPlaceSimpleMemoryAllocator(
error_reporter, aligned_arena, aligned_arena_size);
TfLiteStatus status = Init();
TfLiteStatus status = InitGraphAndContextTensorData();
// TODO(b/147871299): Consider improving this code. A better way of handling
// failures in the constructor is to have a static function that returns a
// pointer to the class. If allocation failed, a nullptr will be returned.
@ -463,88 +415,15 @@ MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model,
}
}
TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations(
TfLiteStatus MicroAllocator::InitializeFromFlatbuffer(
const OpResolver& op_resolver,
NodeAndRegistration** node_and_registrations) {
if (!active_) {
return kTfLiteError;
}
auto* output = reinterpret_cast<NodeAndRegistration*>(
memory_allocator_->AllocateFromTail(
sizeof(NodeAndRegistration) * subgraph_->operators()->size(),
alignof(NodeAndRegistration)));
if (output == nullptr) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Failed to allocate memory for node_and_registrations.");
return kTfLiteError;
}
TfLiteStatus status = kTfLiteOk;
auto* opcodes = model_->operator_codes();
MicroBuiltinDataAllocator builtin_data_allocator(memory_allocator_);
for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
const auto* op = subgraph_->operators()->Get(i);
size_t index = op->opcode_index();
if (index >= opcodes->size()) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Missing registration for opcode_index %d\n", index);
return kTfLiteError;
}
auto* opcode = (*opcodes)[index];
status = GetRegistrationFromOpCode(opcode, op_resolver, error_reporter_,
&(output[i].registration));
if (status != kTfLiteOk) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed to get registration from op code %s\n ",
EnumNameBuiltinOperator(opcode->builtin_code()));
return status;
}
const auto* registration = output[i].registration;
if (registration == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter_, "Skipping op for opcode_index %d\n",
index);
return kTfLiteError;
}
BuiltinOperator op_type =
static_cast<BuiltinOperator>(registration->builtin_code);
if (op_type != BuiltinOperator_CUSTOM && op->custom_options()) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Unsupported behavior: found builtin operator %s with custom "
"options.\n",
EnumNameBuiltinOperator(op_type));
return kTfLiteError;
}
const char* custom_data = nullptr;
size_t custom_data_size = 0;
unsigned char* builtin_data = nullptr;
if (op->custom_options()) {
custom_data = reinterpret_cast<const char*>(op->custom_options()->data());
custom_data_size = op->custom_options()->size();
} else {
TF_LITE_ENSURE_STATUS(ParseOpData(op, op_type, error_reporter_,
&builtin_data_allocator,
(void**)(&builtin_data)));
}
// Disregard const qualifier to workaround with existing API.
TfLiteIntArray* inputs_array = const_cast<TfLiteIntArray*>(
reinterpret_cast<const TfLiteIntArray*>(op->inputs()));
TfLiteIntArray* outputs_array = const_cast<TfLiteIntArray*>(
reinterpret_cast<const TfLiteIntArray*>(op->outputs()));
TfLiteNode* node = &(output[i].node);
*node = {};
node->inputs = inputs_array;
node->outputs = outputs_array;
node->builtin_data = reinterpret_cast<void*>(builtin_data);
node->custom_initial_data = custom_data;
node->custom_initial_data_size = custom_data_size;
}
*node_and_registrations = output;
TF_LITE_ENSURE_STATUS(AllocateNodeAndRegistrations(node_and_registrations));
TF_LITE_ENSURE_STATUS(PrepareNodeAndRegistrationDataFromFlatbuffer(
op_resolver, *node_and_registrations));
return kTfLiteOk;
}
@ -679,4 +558,151 @@ void* MicroAllocator::GetScratchBuffer(int buffer_idx) const {
return scratch_buffer_handles_[scratch_buffer_count_ - buffer_idx - 1].data;
}
size_t MicroAllocator::used_bytes() const {
if (active_) {
return 0;
}
TF_LITE_REPORT_ERROR(error_reporter_, "Total buffer usage: %d bytes",
memory_allocator_->GetUsedBytes());
TF_LITE_REPORT_ERROR(error_reporter_, "Head usage: %d bytes",
memory_allocator_->GetHeadUsedBytes());
TF_LITE_REPORT_ERROR(error_reporter_, "Tail usage: %d bytes",
memory_allocator_->GetTailUsedBytes());
return memory_allocator_->GetUsedBytes();
}
TfLiteStatus MicroAllocator::InitGraphAndContextTensorData() {
auto* subgraphs = model_->subgraphs();
if (subgraphs->size() != 1) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Only 1 subgraph is currently supported.\n");
return kTfLiteError;
}
subgraph_ = (*subgraphs)[0];
TF_LITE_ENSURE_STATUS(AllocateTfLiteTensorArray());
TF_LITE_ENSURE_STATUS(PopulateTfLiteTensorArrayFromFlatbuffer());
return kTfLiteOk;
}
TfLiteStatus MicroAllocator::AllocateTfLiteTensorArray() {
context_->tensors_size = subgraph_->tensors()->size();
context_->tensors =
reinterpret_cast<TfLiteTensor*>(memory_allocator_->AllocateFromTail(
sizeof(TfLiteTensor) * context_->tensors_size,
alignof(TfLiteTensor)));
if (context_->tensors == nullptr) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Failed to allocate memory for context->tensors, %d bytes required",
sizeof(TfLiteTensor) * context_->tensors_size);
return kTfLiteError;
}
return kTfLiteOk;
}
TfLiteStatus MicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer() {
// Initialize tensors in context_ using the flatbuffer for quantization data.
for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) {
TfLiteStatus status = internal::InitializeTfLiteTensorFromFlatbuffer(
memory_allocator_, *subgraph_->tensors()->Get(i), model_->buffers(),
error_reporter_, &context_->tensors[i]);
if (status != kTfLiteOk) {
TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d",
i);
return kTfLiteError;
}
}
return kTfLiteOk;
}
TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations(
NodeAndRegistration** node_and_registrations) {
NodeAndRegistration* output = reinterpret_cast<NodeAndRegistration*>(
memory_allocator_->AllocateFromTail(
sizeof(NodeAndRegistration) * subgraph_->operators()->size(),
alignof(NodeAndRegistration)));
if (output == nullptr) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Failed to allocate memory for node_and_registrations.");
return kTfLiteError;
}
*node_and_registrations = output;
return kTfLiteOk;
}
TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
const OpResolver& op_resolver,
NodeAndRegistration* node_and_registrations) {
TfLiteStatus status = kTfLiteOk;
auto* opcodes = model_->operator_codes();
MicroBuiltinDataAllocator builtin_data_allocator(memory_allocator_);
for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
const auto* op = subgraph_->operators()->Get(i);
const size_t index = op->opcode_index();
if (index >= opcodes->size()) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Missing registration for opcode_index %d\n", index);
return kTfLiteError;
}
auto* opcode = (*opcodes)[index];
status =
GetRegistrationFromOpCode(opcode, op_resolver, error_reporter_,
&(node_and_registrations[i].registration));
if (status != kTfLiteOk) {
TF_LITE_REPORT_ERROR(error_reporter_,
"Failed to get registration from op code %s\n ",
EnumNameBuiltinOperator(opcode->builtin_code()));
return status;
}
const auto* registration = node_and_registrations[i].registration;
if (registration == nullptr) {
TF_LITE_REPORT_ERROR(error_reporter_, "Skipping op for opcode_index %d\n",
index);
return kTfLiteError;
}
BuiltinOperator op_type =
static_cast<BuiltinOperator>(registration->builtin_code);
if (op_type != BuiltinOperator_CUSTOM && op->custom_options()) {
TF_LITE_REPORT_ERROR(
error_reporter_,
"Unsupported behavior: found builtin operator %s with custom "
"options.\n",
EnumNameBuiltinOperator(op_type));
return kTfLiteError;
}
const char* custom_data = nullptr;
size_t custom_data_size = 0;
unsigned char* builtin_data = nullptr;
if (op->custom_options()) {
custom_data = reinterpret_cast<const char*>(op->custom_options()->data());
custom_data_size = op->custom_options()->size();
} else {
TF_LITE_ENSURE_STATUS(ParseOpData(op, op_type, error_reporter_,
&builtin_data_allocator,
(void**)(&builtin_data)));
}
// Disregard const qualifier to workaround with existing API.
TfLiteIntArray* inputs_array = const_cast<TfLiteIntArray*>(
reinterpret_cast<const TfLiteIntArray*>(op->inputs()));
TfLiteIntArray* outputs_array = const_cast<TfLiteIntArray*>(
reinterpret_cast<const TfLiteIntArray*>(op->outputs()));
TfLiteNode* node = &(node_and_registrations[i].node);
*node = {};
node->inputs = inputs_array;
node->outputs = outputs_array;
node->builtin_data = reinterpret_cast<void*>(builtin_data);
node->custom_initial_data = custom_data;
node->custom_initial_data_size = custom_data_size;
}
return kTfLiteOk;
}
} // namespace tflite

58
tensorflow/lite/micro/micro_allocator.h
View File

@ -30,9 +30,9 @@ namespace tflite {
// Namespace used for unittests.
namespace internal {
// Sets up all of the data structure members for a runtime tensor
// based on the contents of a serialized tensor.
TfLiteStatus InitializeRuntimeTensor(
// Sets up all of the data structure members for a TfLiteTensor based on the
// contents of a serialized tensor in the flatbuffer.
TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
SimpleMemoryAllocator* allocator, const tflite::Tensor& flatbuffer_tensor,
const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers,
ErrorReporter* error_reporter, TfLiteTensor* result);
@ -86,6 +86,15 @@ class MicroAllocator {
uint8_t* tensor_arena, size_t arena_size,
ErrorReporter* error_reporter);
// Run through the model flatbuffer data (loaded from the TfLiteModel
// instance) to allocate nodes and registrations. We need to keep them for the
// entire life time of the model to allow persistent tensors. This method
// needs to be called before FinishTensorAllocation method. This method also
// allocates any internal Op data that is required from the flatbuffer.
TfLiteStatus InitializeFromFlatbuffer(
const OpResolver& op_resolver,
NodeAndRegistration** node_and_registrations);
// Runs through the model and allocates all necessary input, output and
// intermediate tensors.
// WARNING: doing any allocation after calling this method has the risk of
@ -93,17 +102,6 @@ class MicroAllocator {
// called in this class.
TfLiteStatus FinishTensorAllocation();
// Returns the arena usage in bytes, only available after
// `FinishTensorAllocation`. Otherwise, it will return 0.
size_t used_bytes() const;
// Run through the model to allocate nodes and registrations. We need to keep
// them for the entire life time of the model to allow persistent tensors.
// This method needs to be called before FinishTensorAllocation method.
TfLiteStatus AllocateNodeAndRegistrations(
const OpResolver& op_resolver,
NodeAndRegistration** node_and_registrations);
// Allocates persistent buffer which has the same life time as the allocator.
// The memory is immediately available and is allocated from the tail of the
// arena.
@ -120,8 +118,38 @@ class MicroAllocator {
// Returns the pointer to the planned scratch buffer.
void* GetScratchBuffer(int buffer_idx) const;
// Returns the arena usage in bytes, only available after
// `FinishTensorAllocation`. Otherwise, it will return 0.
size_t used_bytes() const;
protected:
// Allocates an array in the arena to hold pointers to the tensors required
// to initialize and prepare a model. These allocations are stored and
// populated on the context.
TfLiteStatus AllocateTfLiteTensorArray();
// Populates content on the list of tensor pointers required to initialize and
// prepare a model from data in the flatbuffer (loaded from the TfLiteModel
// instance). Persistent data (e.g. quantization params) is allocated from the
// arena.
TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer();
// Allocates an array in the arena to hold pointers to the node and
// registration pointers required to represent the inference graph of the
// model.
TfLiteStatus AllocateNodeAndRegistrations(
NodeAndRegistration** node_and_registrations);
// Populates node and registration pointers representing the inference graph
// of the model from values inside the flatbuffer (loaded from the TfLiteModel
// instance). Persistent data (e.g. operator data) is allocated from the
// arena.
TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
const OpResolver& op_resolver,
NodeAndRegistration* node_and_registrations);
private:
TfLiteStatus Init();
TfLiteStatus InitGraphAndContextTensorData();
const Model* model_;
// A simple memory allocator that always allocate from the arena tail.

6
tensorflow/lite/micro/micro_allocator_test.cc
View File

@ -77,7 +77,7 @@ TF_LITE_MICRO_TEST(TestInitializeRuntimeTensor) {
TfLiteTensor allocated_tensor;
TF_LITE_MICRO_EXPECT_EQ(
kTfLiteOk, tflite::internal::InitializeRuntimeTensor(
kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer(
&simple_allocator, *tensor, buffers, micro_test::reporter,
&allocated_tensor));
TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type);
@ -103,7 +103,7 @@ TF_LITE_MICRO_TEST(TestInitializeQuantizedTensor) {
TfLiteTensor allocated_tensor;
TF_LITE_MICRO_EXPECT_EQ(
kTfLiteOk, tflite::internal::InitializeRuntimeTensor(
kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer(
&simple_allocator, *tensor, buffers, micro_test::reporter,
&allocated_tensor));
TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type);
@ -129,7 +129,7 @@ TF_LITE_MICRO_TEST(TestMissingQuantization) {
TfLiteTensor allocated_tensor;
TF_LITE_MICRO_EXPECT_EQ(
kTfLiteOk, tflite::internal::InitializeRuntimeTensor(
kTfLiteOk, tflite::internal::InitializeTfLiteTensorFromFlatbuffer(
&simple_allocator, *tensor, buffers, micro_test::reporter,
&allocated_tensor));
TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, allocated_tensor.type);

2
tensorflow/lite/micro/micro_interpreter.cc
View File

@ -165,7 +165,7 @@ void MicroInterpreter::CorrectTensorDataEndianness(T* data, int32_t size) {
}
TfLiteStatus MicroInterpreter::AllocateTensors() {
TF_LITE_ENSURE_OK(&context_, allocator_.AllocateNodeAndRegistrations(
TF_LITE_ENSURE_OK(&context_, allocator_.InitializeFromFlatbuffer(
op_resolver_, &node_and_registrations_));
// Only allow AllocatePersistentBuffer in Init stage.