From 57c7fc537406eb754bffe5f723a21cb1ecb72f14 Mon Sep 17 00:00:00 2001 From: "Kwabena W. Agyeman" Date: Sat, 14 Jun 2025 19:51:06 -0700 Subject: [PATCH] modules/py_ml: Return tensor references for post-processors. Converting the output tensors into floats for the prost-processors causes memory exhaustion when models become very large. Additionally, it wastes processing time converting values which may not be used. By moving the conversion step into the post-processors we avoid this issue. If no callback is passed for post-processing the converted output to a floating point ndarray is returned still. --- modules/py_ml.c | 63 +++++++++++++---------- scripts/libraries/ml/ml/postprocessing.py | 32 +++++++++--- 2 files changed, 61 insertions(+), 34 deletions(-) diff --git a/modules/py_ml.c b/modules/py_ml.c index 72c191ee5..e80c6ccd9 100644 --- a/modules/py_ml.c +++ b/modules/py_ml.c @@ -65,7 +65,7 @@ static size_t py_ml_tuple_sum(mp_obj_tuple_t *o) { return size; } -static size_t pl_ml_dtype_size(char dtype) { +static size_t py_ml_dtype_size(char dtype) { switch (dtype) { case 'f': return 4; @@ -92,7 +92,7 @@ static void py_ml_process_input(py_ml_model_obj_t *model, mp_obj_t arg) { if (mp_obj_is_callable(input_arg)) { // Input is a callable. Call the object and pass the tensor buffer and dtype. mp_obj_t fargs[3] = { - mp_obj_new_bytearray_by_ref(input_size * pl_ml_dtype_size(input_dtype), input_buffer), + mp_obj_new_bytearray_by_ref(input_size * py_ml_dtype_size(input_dtype), input_buffer), MP_OBJ_FROM_PTR(input_shape), mp_obj_new_int(input_dtype) }; @@ -151,7 +151,7 @@ static void py_ml_process_input(py_ml_model_obj_t *model, mp_obj_t arg) { } } -static mp_obj_t py_ml_process_output(py_ml_model_obj_t *model) { +static mp_obj_t py_ml_process_output(py_ml_model_obj_t *model, bool callback) { mp_obj_list_t *output_list = MP_OBJ_TO_PTR(mp_obj_new_list(model->outputs_size, NULL)); for (size_t i = 0; i < model->outputs_size; i++) { void *model_output = ml_backend_get_output(model, i); @@ -172,31 +172,38 @@ static mp_obj_t py_ml_process_output(py_ml_model_obj_t *model) { shape[ulab_offset + j] = mp_obj_get_int(output_shape->items[j]); } - ndarray_obj_t *ndarray = ndarray_new_dense_ndarray(output_shape->len, shape, NDARRAY_FLOAT); + ndarray_obj_t *ndarray; - if (output_dtype == 'f') { - memcpy(ndarray->array, model_output, size * sizeof(float)); - } else if (output_dtype == 'b') { - for (size_t j = 0; j < size; j++) { - float v = (((int8_t *) model_output)[j] - output_zero_point); - ((float *) ndarray->array)[j] = v * output_scale; - } - } else if (output_dtype == 'B') { - for (size_t j = 0; j < size; j++) { - float v = (((uint8_t *) model_output)[j] - output_zero_point); - ((float *) ndarray->array)[j] = v * output_scale; - } - } else if (output_dtype == 'h') { - for (size_t j = 0; j < size; j++) { - float v = (((int16_t *) model_output)[j] - output_zero_point); - ((float *) ndarray->array)[j] = v * output_scale; - } - } else if (output_dtype == 'H') { - for (size_t j = 0; j < size; j++) { - float v = (((uint16_t *) model_output)[j] - output_zero_point); - ((float *) ndarray->array)[j] = v * output_scale; + if (callback) { + ndarray = ndarray_new_ndarray(output_shape->len, shape, NULL, output_dtype, model_output); + } else { + ndarray = ndarray_new_dense_ndarray(output_shape->len, shape, NDARRAY_FLOAT); + + if (output_dtype == 'f') { + memcpy(ndarray->array, model_output, size * sizeof(float)); + } else if (output_dtype == 'b') { + for (size_t j = 0; j < size; j++) { + float v = (((int8_t *) model_output)[j] - output_zero_point); + ((float *) ndarray->array)[j] = v * output_scale; + } + } else if (output_dtype == 'B') { + for (size_t j = 0; j < size; j++) { + float v = (((uint8_t *) model_output)[j] - output_zero_point); + ((float *) ndarray->array)[j] = v * output_scale; + } + } else if (output_dtype == 'h') { + for (size_t j = 0; j < size; j++) { + float v = (((int16_t *) model_output)[j] - output_zero_point); + ((float *) ndarray->array)[j] = v * output_scale; + } + } else if (output_dtype == 'H') { + for (size_t j = 0; j < size; j++) { + float v = (((uint16_t *) model_output)[j] - output_zero_point); + ((float *) ndarray->array)[j] = v * output_scale; + } } } + output_list->items[i] = MP_OBJ_FROM_PTR(ndarray); } @@ -254,6 +261,8 @@ static mp_obj_t py_ml_model_predict(size_t n_args, const mp_obj_t *pos_args, mp_ mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unsupported input type. Expected a list")); } + bool callback = args[ARG_callback].u_obj != mp_const_none; + OMV_PROFILE_START(preprocess); py_ml_process_input(model, pos_args[1]); OMV_PROFILE_PRINT(preprocess); @@ -263,10 +272,10 @@ static mp_obj_t py_ml_model_predict(size_t n_args, const mp_obj_t *pos_args, mp_ OMV_PROFILE_PRINT(inference); OMV_PROFILE_START(postprocess); - mp_obj_t output = py_ml_process_output(model); + mp_obj_t output = py_ml_process_output(model, callback); OMV_PROFILE_PRINT(postprocess); - if (args[ARG_callback].u_obj != mp_const_none) { + if (callback) { // Pass model, inputs, outputs to the post-processing callback. mp_obj_t fargs[3] = { MP_OBJ_FROM_PTR(model), pos_args[1], output }; OMV_PROFILE_START(postprocess_callback); diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index badcc03b5..0b21271b9 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -36,6 +36,12 @@ from ulab import numpy as np _NO_DETECTION = const(()) +def dequantize(value, dtype, zero_point, scale): + if dtype == 'f': + return value + return (value - zero_point) * scale + + # FOMO generates an image per class, where each pixel represents the centroid # of the trained object. These images are processed with `find_blobs()` to # extract centroids, and `get_stats()` is used to get their scores. Overlapping @@ -48,9 +54,12 @@ class fomo_postprocess: def __call__(self, model, inputs, outputs): n, oh, ow, oc = model.output_shape[0] + s = model.output_scale[0] + zp = model.output_zero_point[0] + dt = model.output_dtype[0] nms = NMS(ow, oh, inputs[0].roi) for i in range(oc): - img = image.Image(outputs[0][0, :, :, i] * 255) + img = image.Image(dequantize(outputs[0][0, :, :, i], dt, zp, s) * 255) blobs = img.find_blobs( self.threshold_list, x_stride=1, area_threshold=1, pixels_threshold=1 ) @@ -89,6 +98,9 @@ class yolo_v2_postprocess: def __call__(self, model, inputs, outputs): ob, oh, ow, oc = model.output_shape[0] + s = model.output_scale[0] + zp = model.output_zero_point[0] + dt = model.output_dtype[0] class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES def sigmoid(x): @@ -106,13 +118,13 @@ class yolo_v2_postprocess: _YOLO_V2_CLASSES + class_count)) # Threshold all the scores - score_indices = sigmoid(row_outputs[:, _YOLO_V2_SCORE]) + score_indices = sigmoid(dequantize(row_outputs[:, _YOLO_V2_SCORE], dt, zp, s)) score_indices = np.nonzero(score_indices > self.threshold)[0] if not len(score_indices): return _NO_DETECTION # Get the bounding boxes that have a valid score - bb = np.take(row_outputs, score_indices, axis=0) + bb = dequantize(np.take(row_outputs, score_indices, axis=0), dt, zp, s) # Extract rows, columns, and anchor indices bb_rows = score_indices // (ow * self.anchors_len) @@ -179,19 +191,22 @@ class yolo_v5_postprocess: def __call__(self, model, inputs, outputs): oh, ow, oc = model.output_shape[0] + s = model.output_scale[0] + zp = model.output_zero_point[0] + dt = model.output_dtype[0] class_count = oc - _YOLO_V5_CLASSES # Reshape the output to a 2D array row_outputs = outputs[0].reshape((oh * ow, _YOLO_V5_CLASSES + class_count)) # Threshold all the scores - score_indices = row_outputs[:, _YOLO_V5_SCORE] + score_indices = dequantize(row_outputs[:, _YOLO_V5_SCORE], dt, zp, s) score_indices = np.nonzero(score_indices > self.threshold)[0] if not len(score_indices): return _NO_DETECTION # Get the bounding boxes that have a valid score - bb = np.take(row_outputs, score_indices, axis=0) + bb = dequantize(np.take(row_outputs, score_indices, axis=0), dt, zp, s) # Get the score information bb_scores = bb[:, _YOLO_V5_SCORE] @@ -233,19 +248,22 @@ class yolo_v8_postprocess: def __call__(self, model, inputs, outputs): oh, ow, oc = model.output_shape[0] + s = model.output_scale[0] + zp = model.output_zero_point[0] + dt = model.output_dtype[0] class_count = ow - _YOLO_V8_CLASSES # Reshape the output to a 2D array column_outputs = outputs[0].reshape((oh * (_YOLO_V8_CLASSES + class_count), oc)) # Threshold all the scores - score_indices = np.max(column_outputs[_YOLO_V8_CLASSES:, :], axis=0) + score_indices = np.max(dequantize(column_outputs[_YOLO_V8_CLASSES:, :], dt, zp, s), axis=0) score_indices = np.nonzero(score_indices > self.threshold)[0] if not len(score_indices): return _NO_DETECTION # Get the bounding boxes that have a valid score - bb = np.take(column_outputs, score_indices, axis=1) + bb = dequantize(np.take(column_outputs, score_indices, axis=1), dt, zp, s) # Get the score information bb_scores = np.max(bb[_YOLO_V8_CLASSES:, :], axis=0)