From 3564081ab864b9dcc0a21fecbb9427c186224a10 Mon Sep 17 00:00:00 2001 From: Tiezhen WANG Date: Wed, 6 May 2020 18:43:14 -0700 Subject: [PATCH] TFLM: Print out the buffer usage. Head space is reusable while the tail space is persistent. This gives some guidance on how much ram could be saved by using multi-tenant TFLM. Check test `TestFinishTensorAllocation` in micro_allocator_test.cc for the usage. PiperOrigin-RevId: 310271606 Change-Id: I5ad75fb4a01504ba584d2af036f474869270bba1 --- tensorflow/lite/micro/micro_allocator.cc | 13 +++++++++++++ tensorflow/lite/micro/micro_allocator.h | 7 +------ tensorflow/lite/micro/simple_memory_allocator.h | 3 +++ 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 28de77cdb6c..c2b32166467 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -419,6 +419,19 @@ TfLiteStatus MicroAllocator::Init() { 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) diff --git a/tensorflow/lite/micro/micro_allocator.h b/tensorflow/lite/micro/micro_allocator.h index a846b0c63ba..b16f814071c 100644 --- a/tensorflow/lite/micro/micro_allocator.h +++ b/tensorflow/lite/micro/micro_allocator.h @@ -91,12 +91,7 @@ class MicroAllocator { // Returns the arena usage in bytes, only available after // `FinishTensorAllocation`. Otherwise, it will return 0. - size_t used_bytes() const { - if (active_) { - return 0; - } - return memory_allocator_->GetUsedBytes(); - } + 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. diff --git a/tensorflow/lite/micro/simple_memory_allocator.h b/tensorflow/lite/micro/simple_memory_allocator.h index cf1818609f6..ed73104a2c6 100644 --- a/tensorflow/lite/micro/simple_memory_allocator.h +++ b/tensorflow/lite/micro/simple_memory_allocator.h @@ -51,6 +51,9 @@ class SimpleMemoryAllocator { size_t GetAvailableMemory() const { return tail_ - head_; } size_t GetUsedBytes() const { return GetBufferSize() - GetAvailableMemory(); } + size_t GetHeadUsedBytes() const { return head_ - buffer_head_; } + size_t GetTailUsedBytes() const { return buffer_tail_ - tail_; } + private: size_t GetBufferSize() const { return buffer_tail_ - buffer_head_; }