[SDK] compatible with tflite 2.1.0

see commit 8599f7c [TFLite] Model importer to be compatible with tflite 2.1.0 (apache#5497) Signed-off-by: Taikang Hu <[email protected]>
windclarion · Jul 2, 2020 · ecfa27d · ecfa27d
1 parent 03fa40b
commit ecfa27d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/python/asr/nn.py b/python/asr/nn.py
@@ -74,10 +74,16 @@ def _load_frontend_model(path, input_tensor):
         input_dtype = 'float32'
 
     if info.frontend == 'tflite':
-        import tflite.Model
+        import tflite
 
         tflite_model_buf = open(path, 'rb').read()
-        tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
+        # TFLite.Model.Model has changed to TFLite.Model from 1.14 to 2.1
+        try:
+            tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
+        except AttributeError:
+            tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
+        except ImportError:
+            raise ImportError("The tflite package must be installed")
 
         mod, param = relay.frontend.from_tflite(tflite_model,
                                                  shape_dict={input_tensor: input_shape},