cnn: Try fix hardcoded output handling

Would give corrupt results for anything else than outputs with 10 lenght
The .run() function now requires an array.array('f') where to place output
There is a new method, output_dimensions() to get the shape of the output

Author: jonnor

src/tinymaix_cnn/mod_cnn.c

@@ -34,39 +34,30 @@ void NORETURN abort() {
 }
 #endif
 
+// Global variable to capture data from layer_cb
+uint16_t *g_out_dims = NULL;
 
+// XXX: caller must set g_out_dims
 static tm_err_t layer_cb(tm_mdl_t* mdl, tml_head_t* lh)
-{  
-#if 0
-    //dump middle result
-    int h = lh->out_dims[1];
-    int w = lh->out_dims[2];
-    int ch= lh->out_dims[3];
-    mtype_t* output = TML_GET_OUTPUT(mdl, lh);
-    return TM_OK;
-    TM_PRINTF("Layer %d callback ========\n", mdl->layer_i);
-    #if 1
-    for(int y=0; y<h; y++){
-        TM_PRINTF("[");
-        for(int x=0; x<w; x++){
-            TM_PRINTF("[");
-            for(int c=0; c<ch; c++){
-            #if TM_MDL_TYPE == TM_MDL_FP32
-                TM_PRINTF("%.3f,", output[(y*w+x)*ch+c]);
-            #else
-                TM_PRINTF("%.3f,", TML_DEQUANT(lh,output[(y*w+x)*ch+c]));
-            #endif
-            }
-            TM_PRINTF("],");
+{
+    const int h = lh->out_dims[1];
+    const int w = lh->out_dims[2];
+    const int ch= lh->out_dims[3];
+    const bool is_out = lh->is_out;
+
+#if DEBUG
+    mp_printf(&mp_plat_print,
+        "cnn-layer-cb is_out=%d h=%d w=%d ch=%d size=%d \n",
+        (int)is_out, h, w, ch);
+#endif
+
+    if (is_out) {
+        for (int i=0; i<4; i++) {
+            g_out_dims[i] = lh->out_dims[i];
         }
-        TM_PRINTF("],\n");
+
     }
-    TM_PRINTF("\n");
-    #endif
-    return TM_OK;
-#else
     return TM_OK;
-#endif
 }
 
 #define DEBUG (1)
@@ -79,6 +70,7 @@ typedef struct _mp_obj_mod_cnn_t {
     tm_mat_t input;
     uint8_t *model_buffer;
     uint8_t *data_buffer;
+    uint16_t out_dims[4];
 } mp_obj_mod_cnn_t;
 
 mp_obj_full_type_t mod_cnn_type;
@@ -116,11 +108,17 @@ static mp_obj_t mod_cnn_new(mp_obj_t model_data_obj) {
 
     // loading model
     // will set the dimensions of the input matrix
+    o->out_dims[0] = 0;
+    g_out_dims = o->out_dims;
     tm_err_t load_err = tm_load(model, o->model_buffer, o->data_buffer, layer_cb, &o->input);
     if (load_err != TM_OK) {
         mp_raise_msg(&mp_type_RuntimeError, MP_ERROR_TEXT("tm_load error"));
     }
 
+#if DEBUG
+    mp_printf(&mp_plat_print, "cnn-new-done out.dims=%d \n", o->out_dims[0]);
+#endif
+
     return MP_OBJ_FROM_PTR(o);
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(mod_cnn_new_obj, mod_cnn_new);
@@ -141,15 +139,15 @@ static MP_DEFINE_CONST_FUN_OBJ_1(mod_cnn_del_obj, mod_cnn_del);
 
 
 // Add a node to the tree
-static mp_obj_t mod_cnn_run(mp_obj_t self_obj, mp_obj_t input_obj) {
+static mp_obj_t mod_cnn_run(mp_obj_t self_obj, mp_obj_t input_obj, mp_obj_t output_obj) {
 
     mp_obj_mod_cnn_t *o = MP_OBJ_TO_PTR(self_obj);
 
     // Extract input
     mp_buffer_info_t bufinfo;
     mp_get_buffer_raise(input_obj, &bufinfo, MP_BUFFER_RW);
     if (bufinfo.typecode != 'B') {
-        mp_raise_ValueError(MP_ERROR_TEXT("expecting float array"));
+        mp_raise_ValueError(MP_ERROR_TEXT("expecting byte array"));
     }
     uint8_t *input_buffer = bufinfo.buf;
     const int input_length = bufinfo.len / sizeof(*input_buffer);
@@ -160,6 +158,21 @@ static mp_obj_t mod_cnn_run(mp_obj_t self_obj, mp_obj_t input_obj) {
         mp_raise_ValueError(MP_ERROR_TEXT("wrong input size"));
     }
 
+    // Extract output
+    mp_get_buffer_raise(output_obj, &bufinfo, MP_BUFFER_RW);
+    if (bufinfo.typecode != 'f') {
+        mp_raise_ValueError(MP_ERROR_TEXT("expecting float array"));
+    }
+    float *output_buffer = bufinfo.buf;
+    const int output_length = bufinfo.len / sizeof(*output_buffer);
+
+
+    // check buffer size wrt input
+    const int expect_out_length = o->out_dims[1]*o->out_dims[2]*o->out_dims[3];
+    if (output_length != expect_out_length) {
+        mp_raise_ValueError(MP_ERROR_TEXT("wrong output size"));
+    }
+
     // Preprocess data
     tm_mat_t in_uint8 = o->input;
     in_uint8.data = (mtype_t*)input_buffer;
@@ -181,27 +194,33 @@ static mp_obj_t mod_cnn_run(mp_obj_t self_obj, mp_obj_t input_obj) {
         mp_raise_ValueError(MP_ERROR_TEXT("run error"));
     }
 
+    // Copy output into
     tm_mat_t out = outs[0];
-    float* data  = out.dataf;
-    float maxp = 0;
-    int maxi = -1;
-
-    // TODO: pass the entire output vector out to Python
-    // FIXME: unhardcode output handling
-    for(int i=0; i<10; i++){
-        //printf("%d: %.3f\n", i, data[i]);
-        if (data[i] > maxp) {
-            maxi = i;
-            maxp = data[i];
-        }
+    for(int i=0; i<expect_out_length; i++){
+        output_buffer[i] = out.dataf[i];
     }
 
-    return mp_obj_new_int(maxi);
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_3(mod_cnn_run_obj, mod_cnn_run);
+
+
+// Return the shape of the output
+static mp_obj_t mod_cnn_output_dimensions(mp_obj_t self_obj) {
+
+    mp_obj_mod_cnn_t *o = MP_OBJ_TO_PTR(self_obj);
+    const int dimensions = o->out_dims[0];
+    mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(dimensions, NULL));
+
+    for (int i=0; i<dimensions; i++) {
+        tuple->items[i] = mp_obj_new_int(o->out_dims[i+1]);
+    }
+    return tuple;
 }
-static MP_DEFINE_CONST_FUN_OBJ_2(mod_cnn_run_obj, mod_cnn_run);
+static MP_DEFINE_CONST_FUN_OBJ_1(mod_cnn_output_dimensions_obj, mod_cnn_output_dimensions);
 
 
-mp_map_elem_t mod_locals_dict_table[2];
+mp_map_elem_t mod_locals_dict_table[3];
 static MP_DEFINE_CONST_DICT(mod_locals_dict, mod_locals_dict_table);
 
 // This is the entry point and is called when the module is imported
@@ -217,6 +236,7 @@ mp_obj_t mpy_init(mp_obj_fun_bc_t *self, size_t n_args, size_t n_kw, mp_obj_t *a
     // methods
     mod_locals_dict_table[0] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_run), MP_OBJ_FROM_PTR(&mod_cnn_run_obj) };
     mod_locals_dict_table[1] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR___del__), MP_OBJ_FROM_PTR(&mod_cnn_del_obj) };
+    mod_locals_dict_table[2] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_output_dimensions), MP_OBJ_FROM_PTR(&mod_cnn_output_dimensions_obj) };
 
     MP_OBJ_TYPE_SET_SLOT(&mod_cnn_type, locals_dict, (void*)&mod_locals_dict, 2);
 

tests/test_cnn.py

@@ -12,6 +12,9 @@ def test_cnn_create():
         model_data = array.array('B', f.read())
         model = emlearn_cnn.new(model_data)
 
+        out_shape = model.output_dimensions()
+        assert out_shape == (10,)
+
         # TODO: enable these checks
         #wrong_type = array.array('f', [])
         #model.run(wrong_type)