Merge pull request #2403 from openmv/tflm_fix

lib/tflm: Use GC's free memory for the temporary tensor arena.
This commit is contained in:
Ibrahim Abdelkader 2024-09-04 09:55:36 +03:00 committed by GitHub
commit 5fe32bdc8a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -25,15 +25,16 @@ extern "C" {
#include "py/objlist.h" #include "py/objlist.h"
#include "py/objtuple.h" #include "py/objtuple.h"
#include "py/binary.h" #include "py/binary.h"
#include "py/gc.h"
#include "py_ml.h" #include "py_ml.h"
#include "fb_alloc.h"
using namespace tflite; using namespace tflite;
#define TF_ARENA_ALIGNMENT (16 - 1) #define TF_ARENA_ALIGN (16 - 1)
#define TF_ARENA_ROUND(x) (((x) + TF_ARENA_ALIGN) & ~(TF_ARENA_ALIGN))
typedef MicroMutableOpResolver<113> MicroOpsResolver; typedef MicroMutableOpResolver<113> MicroOpsResolver;
typedef struct ml_backend_state { typedef struct ml_backend_state {
void *arena; uint8_t *arena;
const Model *model; const Model *model;
MicroOpsResolver *resolver; MicroOpsResolver *resolver;
MicroInterpreter *interpreter; MicroInterpreter *interpreter;
@ -192,38 +193,57 @@ static void ml_backend_init_ops_resolver(MicroOpsResolver *resolver) {
int ml_backend_init_model(py_ml_model_obj_t *model) { int ml_backend_init_model(py_ml_model_obj_t *model) {
RegisterDebugLogCallback(ml_backend_log_handler); RegisterDebugLogCallback(ml_backend_log_handler);
// Parse model's data. // Parse the model's data.
const Model *tflite_model = GetModel(model->data); const Model *tflite_model = GetModel(model->data);
if (tflite_model->version() != TFLITE_SCHEMA_VERSION) { if (tflite_model->version() != TFLITE_SCHEMA_VERSION) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported model schema")); mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported model schema"));
} }
// Initialize the op resolver. // Initialize a temporary op resolver.
MicroOpsResolver resolver; MicroOpsResolver resolver;
ml_backend_init_ops_resolver(&resolver); ml_backend_init_ops_resolver(&resolver);
// Allocate the interpreter and tensors once to initialize the model, check input gc_info_t info;
// and output data types and to get the optimal tensor arena size. gc_info(&info);
fb_alloc_mark(); // Allocate a temporary interpreter to get the optimal arena size.
uint32_t tensor_arena_size; size_t arena_size = info.max_free * MICROPY_BYTES_PER_GC_BLOCK;
uint8_t *tensor_arena = (uint8_t *) fb_alloc_all(&tensor_arena_size, FB_ALLOC_PREFER_SIZE | FB_ALLOC_CACHE_ALIGN); uint8_t *arena_memory = m_new(uint8_t, arena_size);
MicroInterpreter interpreter(tflite_model, resolver, arena_memory, arena_size);
MicroInterpreter interpreter(tflite_model,
resolver,
tensor_arena,
tensor_arena_size);
if (interpreter.AllocateTensors() != kTfLiteOk) { if (interpreter.AllocateTensors() != kTfLiteOk) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Failed to allocate tensors")); mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Failed to allocate tensors"));
} }
// Round up the optimal arena size to a multiple of the alignment.
arena_size = TF_ARENA_ROUND(interpreter.arena_used_bytes());
m_free(arena_memory);
model->inputs_size = interpreter.inputs_size(); // Allocate the persistent model state and interpreter.
ml_backend_state_t *state = m_new0(ml_backend_state_t, 1);
state->model = GetModel(model->data);
state->arena = m_new(uint8_t, arena_size);
state->resolver = new(m_new0(MicroOpsResolver, 1)) MicroOpsResolver();
ml_backend_init_ops_resolver(state->resolver);
state->interpreter = new(m_new0(MicroInterpreter, 1)) MicroInterpreter(state->model,
*state->resolver,
state->arena,
arena_size);
if (state->interpreter->AllocateTensors() != kTfLiteOk) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Failed to allocate tensors"));
}
// Initialize the model's state.
model->state = state;
model->memory_addr = (uint32_t) state->arena;
model->memory_size = arena_size;
// Initialize the model's inputs.
model->inputs_size = state->interpreter->inputs_size();
model->input_shape = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL)); model->input_shape = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL));
model->input_scale = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL)); model->input_scale = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL));
model->input_zero_point = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL)); model->input_zero_point = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL));
model->input_dtype = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL)); model->input_dtype = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->inputs_size, NULL));
for (size_t i=0; i<model->inputs_size; i++) { for (size_t i=0; i<model->inputs_size; i++) {
TfLiteTensor *input = interpreter.input(i); TfLiteTensor *input = state->interpreter->input(i);
// Check input data type. // Check input data type.
if (!ml_backend_valid_dataype(input->type)) { if (!ml_backend_valid_dataype(input->type)) {
@ -242,14 +262,15 @@ int ml_backend_init_model(py_ml_model_obj_t *model) {
model->input_dtype->items[i] = mp_obj_new_int(ml_backend_map_dtype(input->type)); model->input_dtype->items[i] = mp_obj_new_int(ml_backend_map_dtype(input->type));
} }
model->outputs_size = interpreter.outputs_size(); // Initialize the model's outputs.
model->outputs_size = state->interpreter->outputs_size();
model->output_shape = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL)); model->output_shape = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL));
model->output_scale = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL)); model->output_scale = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL));
model->output_zero_point = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL)); model->output_zero_point = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL));
model->output_dtype = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL)); model->output_dtype = (mp_obj_tuple_t *) MP_OBJ_TO_PTR(mp_obj_new_tuple(model->outputs_size, NULL));
for (size_t i=0; i<model->outputs_size; i++) { for (size_t i=0; i<model->outputs_size; i++) {
TfLiteTensor *output = interpreter.output(i); TfLiteTensor *output = state->interpreter->output(i);
// Check output data type. // Check output data type.
if (!ml_backend_valid_dataype(output->type)) { if (!ml_backend_valid_dataype(output->type)) {
@ -268,28 +289,6 @@ int ml_backend_init_model(py_ml_model_obj_t *model) {
model->output_dtype->items[i] = mp_obj_new_int(ml_backend_map_dtype(output->type)); model->output_dtype->items[i] = mp_obj_new_int(ml_backend_map_dtype(output->type));
} }
model->memory_size = interpreter.arena_used_bytes() + 1024;
// Free the temporary arena.
fb_alloc_free_till_mark();
// Allocate the persistent state.
ml_backend_state_t *state = m_new0(ml_backend_state_t, 1);
state->model = GetModel(model->data);
state->arena = m_new(char, model->memory_size + TF_ARENA_ALIGNMENT);
state->resolver = new(m_new0(MicroOpsResolver, 1)) MicroOpsResolver();
ml_backend_init_ops_resolver(state->resolver);
uint8_t *aligned_arena = (uint8_t *) (((uintptr_t) state->arena + TF_ARENA_ALIGNMENT) & ~(TF_ARENA_ALIGNMENT));
state->interpreter = new(m_new0(MicroInterpreter, 1)) MicroInterpreter(state->model,
*state->resolver,
aligned_arena,
model->memory_size);
if (state->interpreter->AllocateTensors() != kTfLiteOk) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Failed to allocate tensors"));
}
model->state = state;
model->memory_addr = (uint32_t) state->arena;
return 0; return 0;
} }