Fix TfLite issue (#701)

Co-authored-by: alonre24 <alonreshef24@gmail.com>

Author: lantiga

src/backends/libtflite_c/tflite_c.cpp

@@ -4,6 +4,7 @@
 #include "redismodule.h"
 #include "tensorflow/lite/model.h"
 #include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/util.h"
 #include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/tools/evaluation/utils.h"
 
@@ -52,6 +53,7 @@ static DLDataType getDLDataType(const TfLiteTensor *tensor) {
     return dtype;
 }
 
+
 static DLDevice getDLDevice(const TfLiteTensor *tensor, const int64_t &device_id) {
     DLDevice device;
     device.device_id = device_id;
@@ -75,29 +77,52 @@ size_t dltensorBytes(DLManagedTensor *t) {
 void copyToTfLiteTensor(std::shared_ptr<tflite::Interpreter> interpreter, int tflite_input,
                         DLManagedTensor *input) {
     TfLiteTensor *tensor = interpreter->tensor(tflite_input);
-
     size_t nbytes = dltensorBytes(input);
+    DLDataType dltensor_type = input->dl_tensor.dtype;
+    const char *type_mismatch_msg = "Input tensor type doesn't match the type expected"
+                            " by the model definition";
 
     switch (tensor->type) {
     case kTfLiteUInt8:
-        memcpy(interpreter->typed_tensor<uint8_t>(tflite_input), input->dl_tensor.data, nbytes);
-        break;
+       if (dltensor_type.code != kDLUInt || dltensor_type.bits != 8) {
+           throw std::logic_error(type_mismatch_msg);
+       }
+       memcpy(interpreter->typed_tensor<uint8_t>(tflite_input), input->dl_tensor.data, nbytes);
+       break;
     case kTfLiteInt64:
+        if (dltensor_type.code != kDLInt || dltensor_type.bits != 64) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<int64_t>(tflite_input), input->dl_tensor.data, nbytes);
         break;
     case kTfLiteInt32:
+        if (dltensor_type.code != kDLInt || dltensor_type.bits != 32) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<int32_t>(tflite_input), input->dl_tensor.data, nbytes);
         break;
     case kTfLiteInt16:
+        if (dltensor_type.code != kDLInt || dltensor_type.bits != 16) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<int16_t>(tflite_input), input->dl_tensor.data, nbytes);
         break;
     case kTfLiteInt8:
+        if (dltensor_type.code != kDLInt || dltensor_type.bits != 8) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<int8_t>(tflite_input), input->dl_tensor.data, nbytes);
         break;
     case kTfLiteFloat32:
+        if (dltensor_type.code != kDLFloat || dltensor_type.bits != 32) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<float>(tflite_input), input->dl_tensor.data, nbytes);
         break;
     case kTfLiteBool:
+        if (dltensor_type.code != kDLBool || dltensor_type.bits != 8) {
+            throw std::logic_error(type_mismatch_msg);
+        }
         memcpy(interpreter->typed_tensor<bool>(tflite_input), input->dl_tensor.data, nbytes);
     case kTfLiteFloat16:
         throw std::logic_error("Float16 not currently supported as input tensor data type");
@@ -318,6 +343,38 @@ extern "C" void tfliteRunModel(void *ctx, long n_inputs, DLManagedTensor **input
         return;
     }
 
+    // NOTE: TFLITE requires all tensors in the graph to be explicitly
+    // preallocated before input tensors are memcopied. These are cached
+    // in the session, so we need to check if for instance the batch size
+    // has changed or the shape has changed in general compared to the
+    // previous run and in that case we resize input tensors and call
+    // the AllocateTensor function manually.
+    bool need_reallocation = false;
+    std::vector<int> dims;
+    for (size_t i = 0; i < tflite_inputs.size(); i++) {
+        const TfLiteTensor* tflite_tensor = interpreter->tensor(tflite_inputs[i]);
+        int64_t ndim = inputs[i]->dl_tensor.ndim;
+        int64_t *shape = inputs[i]->dl_tensor.shape;
+        dims.resize(ndim);
+        for (size_t j=0; j < ndim; j++) {
+            dims[j] = shape[j];
+        }
+        if (!tflite::EqualArrayAndTfLiteIntArray(tflite_tensor->dims, dims.size(), dims.data())) {
+            if (interpreter->ResizeInputTensor(i, dims) != kTfLiteOk) {
+                _setError("Failed to resize input tensors", error);
+                return;
+            }
+            need_reallocation = true;
+        }
+    }
+
+    if (need_reallocation) {
+        if (interpreter->AllocateTensors() != kTfLiteOk) {
+            _setError("Failed to allocate tensors", error);
+            return;
+        }
+    }
+
     try {
         for (size_t i = 0; i < tflite_inputs.size(); i++) {
             copyToTfLiteTensor(interpreter, tflite_inputs[i], inputs[i]);

src/redisai.c

@@ -201,10 +201,6 @@ int RedisAI_ModelStore_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **arg
 
     unsigned long long batchsize = 0;
     if (AC_AdvanceIfMatch(&ac, "BATCHSIZE")) {
-        if (backend == RAI_BACKEND_TFLITE) {
-            return RedisModule_ReplyWithError(
-                ctx, "ERR Auto-batching not supported by the TFLITE backend");
-        }
         if (AC_GetUnsignedLongLong(&ac, &batchsize, 0) != AC_OK) {
             return RedisModule_ReplyWithError(ctx, "ERR Invalid argument for BATCHSIZE");
         }

tests/flow/test_data/lite-model_imagenet_mobilenet_v3_small_100_224_classification_5_default_1.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9df60af7ab24d287a54668e845ea7da1c854086b828a4a4cf46c55c403095053
+size 10209756

tests/flow/tests_tensorflow.py

@@ -646,7 +646,6 @@ def test_tensorflow_modelexecute_script_execute_resnet(env):
     inputvar = 'images'
     outputvar = 'output'
 
-
     model_pb, script, labels, img = load_resnet_test_data()
 
     ret = con.execute_command('AI.MODELSTORE', model_name, 'TF', DEVICE,

tests/flow/tests_tflite.py

@@ -1,4 +1,4 @@
-import redis
+import numpy as np
 
 from includes import *
 
@@ -47,7 +47,61 @@ def test_run_tflite_model(env):
     env.assertEqual(values[0], 1)
 
 
-def test_run_tflite_model_errors(env):
+def test_run_tflite_model_autobatch(env):
+    if not TEST_TFLITE:
+        env.debugPrint("skipping {} since TEST_TFLITE=0".format(sys._getframe().f_code.co_name), force=True)
+        return
+
+    con = env.getConnection()
+    model_pb = load_file_content('lite-model_imagenet_mobilenet_v3_small_100_224_classification_5_default_1.tflite')
+    _, _, _, img = load_resnet_test_data()
+    img = img.astype(np.float32) / 255
+
+    ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'TFLITE', 'CPU',
+                              'BATCHSIZE', 4, 'MINBATCHSIZE', 2,
+                              'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+
+    ret = con.execute_command('AI.MODELGET', 'm{1}', 'META')
+    env.assertEqual(len(ret), 16)
+    if DEVICE == "CPU":
+        env.assertEqual(ret[1], b'TFLITE')
+        env.assertEqual(ret[3], b'CPU')
+
+    ret = con.execute_command('AI.TENSORSET', 'a{1}',
+                              'FLOAT', 1, img.shape[1], img.shape[0], 3,
+                              'BLOB', img.tobytes())
+    env.assertEqual(ret, b'OK')
+
+    ret = con.execute_command('AI.TENSORSET', 'b{1}',
+                              'FLOAT', 1, img.shape[1], img.shape[0], 3,
+                              'BLOB', img.tobytes())
+    env.assertEqual(ret, b'OK')
+
+    def run():
+        con = env.getConnection()
+        con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1,
+                            'b{1}', 'OUTPUTS', 1, 'd{1}')
+        ensureSlaveSynced(con, env)
+
+    t = threading.Thread(target=run)
+    t.start()
+
+    con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'a{1}', 'OUTPUTS', 1, 'c{1}')
+    t.join()
+
+    ensureSlaveSynced(con, env)
+
+    values = con.execute_command('AI.TENSORGET', 'c{1}', 'VALUES')
+    idx = np.argmax(values)
+    env.assertEqual(idx, 112)
+
+    values = con.execute_command('AI.TENSORGET', 'd{1}', 'VALUES')
+    idx = np.argmax(values)
+    env.assertEqual(idx, 112)
+
+
+def test_run_tflite_errors(env):
     if not TEST_TFLITE:
         env.debugPrint("skipping {} since TEST_TFLITE=0".format(sys._getframe().f_code.co_name), force=True)
         return
@@ -64,13 +118,6 @@ def test_run_tflite_model_errors(env):
     check_error_message(env, con, "Failed to load model from buffer",
                         'AI.MODELSTORE', 'm{1}', 'TFLITE', 'CPU', 'TAG', 'asdf', 'BLOB', wrong_model_pb)
 
-    # TODO: Autobatch is tricky with TFLITE because TFLITE expects a fixed batch
-    #       size. At least we should constrain MINBATCHSIZE according to the
-    #       hard-coded dims in the tflite model.
-    check_error_message(env, con, "Auto-batching not supported by the TFLITE backend",
-                        'AI.MODELSTORE', 'm{1}', 'TFLITE', 'CPU',
-                        'BATCHSIZE', 2, 'MINBATCHSIZE', 2, 'BLOB', model_pb)
-
     ret = con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw)
     env.assertEqual(ret, b'OK')
 
@@ -82,6 +129,19 @@ def test_run_tflite_model_errors(env):
     check_error_message(env, con, "Number of keys given as INPUTS here does not match model definition",
                         'AI.MODELEXECUTE', 'm_2{1}', 'INPUTS', 3, 'a{1}', 'b{1}', 'c{1}', 'OUTPUTS', 1, 'd{1}')
 
+    model_pb = load_file_content('lite-model_imagenet_mobilenet_v3_small_100_224_classification_5_default_1.tflite')
+    _, _, _, img = load_resnet_test_data()
+
+    ret = con.execute_command('AI.MODELSTORE', 'image_net{1}', 'TFLITE', 'CPU', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ret = con.execute_command('AI.TENSORSET', 'dog{1}', 'UINT8', 1, img.shape[1], img.shape[0], 3,
+                              'BLOB', img.tobytes())
+    env.assertEqual(ret, b'OK')
+
+    # The model expects FLOAT input, but UINT8 tensor is given.
+    check_error_message(env, con, "Input tensor type doesn't match the type expected by the model definition",
+                        'AI.MODELEXECUTE', 'image_net{1}', 'INPUTS', 1, 'dog{1}', 'OUTPUTS', 1, 'output{1}')
+
 
 def test_tflite_modelinfo(env):
     if not TEST_TFLITE: