Change tflite tests from sharkimporter -> sharkdownloader (huggingface#182)

* Change tflite test from sharkimporter -> sharkdownloader

* xfail all uint/int tflite sharkdownloader tests

Author: AmosLewis

shark/shark_downloader.py

@@ -23,6 +23,8 @@
     "bool": np.bool_,
     "int32": np.int32,
     "int64": np.int64,
+    "uint8": np.uint8,
+    "int8": np.int8,
 }
 
 
@@ -32,7 +34,7 @@ def __init__(
         model_name: str,
         tank_url: str = "https://storage.googleapis.com/shark_tank",
         local_tank_dir: str = "./../gen_shark_tank/tflite",
-        model_type: str = "tflite-tosa",
+        model_type: str = "tflite",
         input_json: str = "input.json",
         input_type: str = "int32",
     ):
@@ -84,7 +86,7 @@ def get_inputs(self):
 
     def load_json_input(self):
         print("load json inputs")
-        if self.model_type in ["tflite-tosa"]:
+        if self.model_type in ["tflite"]:
             input_url = (
                 self.tank_url + "/" + str(self.model_name) + "/" + "input.json"
             )
@@ -109,7 +111,7 @@ def load_json_input(self):
         return self.inputs
 
     def load_mlir_model(self):
-        if self.model_type in ["tflite-tosa"]:
+        if self.model_type in ["tflite"]:
             self.mlir_url = (
                 self.tank_url
                 + "/"

tank/person_detect/person_detect_tflite_test.py renamed to tank/albert_lite_base/albert_lite_base_tflite_importer_test.py

@@ -4,30 +4,36 @@
 import pytest
 import unittest
 from shark.parser import shark_args
-import os
-import sys
-import urllib.request
-from PIL import Image
 from shark.tflite_utils import TFLitePreprocessor
 
 
-# model_path = "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea20c54e86d41e3a38312546/tensorflow/lite/micro/models/person_detect.tflite"
-
+# model_path = "https://tfhub.dev/tensorflow/lite-model/albert_lite_base/squadv1/1?lite-format=tflite"
+# model_path = model_path
 
+# Inputs modified to be useful albert inputs.
 def generate_inputs(input_details):
-    exe_basename = os.path.basename(sys.argv[0])
-    workdir = os.path.join(os.path.dirname(__file__), "../tmp", exe_basename)
-    os.makedirs(workdir, exist_ok=True)
-
-    img_path = "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea20c54e86d41e3a38312546/tensorflow/lite/micro/examples/person_detection/testdata/person.bmp"
-    local_path = "/".join([workdir, "person.bmp"])
-    urllib.request.urlretrieve(img_path, local_path)
-
-    shape = input_details[0]["shape"]
-    im = np.array(Image.open(local_path).resize((shape[1], shape[2]))).astype(
-        input_details[0]["dtype"]
+    for input in input_details:
+        print(str(input["shape"]), input["dtype"].__name__)
+
+    args = []
+    args.append(
+        np.random.randint(
+            low=0,
+            high=256,
+            size=input_details[0]["shape"],
+            dtype=input_details[0]["dtype"],
+        )
+    )
+    args.append(
+        np.ones(
+            shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
+        )
+    )
+    args.append(
+        np.zeros(
+            shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
+        )
     )
-    args = [im.reshape(shape)]
     return args
 
 
@@ -41,12 +47,14 @@ def compare_results(mlir_results, tflite_results, details):
         tflite_result = tflite_results[i]
         mlir_result = mlir_result.astype(np.single)
         tflite_result = tflite_result.astype(np.single)
+        print("mlir_result.shape", mlir_result.shape)
+        print("tflite_result.shape", tflite_result.shape)
         assert mlir_result.shape == tflite_result.shape, "shape doesnot match"
         max_error = np.max(np.abs(mlir_result - tflite_result))
         print("Max error (%d): %f", i, max_error)
 
 
-class PersonDetectionTfliteModuleTester:
+class AlbertTfliteModuleTester:
     def __init__(
         self,
         dynamic=False,
@@ -64,25 +72,7 @@ def create_and_check_module(self):
         shark_args.save_vmfb = self.save_vmfb
 
         # Preprocess to get SharkImporter input args
-        # The input has known expected values. We hardcode this value.
-        input_details = [
-            {
-                "shape": [1, 96, 96, 1],
-                "dtype": np.int8,
-                "index": 0,
-            }
-        ]
-        output_details = [
-            {
-                "shape": [1, 2],
-                "dtype": np.int8,
-            }
-        ]
-        tflite_preprocessor = TFLitePreprocessor(
-            model_name="person_detect",
-            input_details=input_details,
-            output_details=output_details,
-        )
+        tflite_preprocessor = TFLitePreprocessor(model_name="albert_lite_base")
         raw_model_file_path = tflite_preprocessor.get_raw_model_file()
         inputs = tflite_preprocessor.get_inputs()
         tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -104,8 +94,20 @@ def create_and_check_module(self):
             mlir_dialect="tflite",
         )
 
-        # Case2: Use manually set inputs
+        # Case1: Use shark_importer default generate inputs
+        shark_module.compile()
+        mlir_results = shark_module.forward(inputs)
+        ## post process results for compare
+        input_details, output_details = tflite_preprocessor.get_model_details()
+        mlir_results = list(mlir_results)
+        for i in range(len(output_details)):
+            dtype = output_details[i]["dtype"]
+            mlir_results[i] = mlir_results[i].astype(dtype)
+        tflite_results = tflite_preprocessor.get_raw_model_output()
+        compare_results(mlir_results, tflite_results, output_details)
 
+        # Case2: Use manually set inputs
+        input_details, output_details = tflite_preprocessor.get_model_details()
         inputs = generate_inputs(input_details)  # new inputs
 
         shark_module = SharkInference(
@@ -117,31 +119,34 @@ def create_and_check_module(self):
         shark_module.compile()
         mlir_results = shark_module.forward(inputs)
         ## post process results for compare
-        # The input has known expected values. We hardcode this value.
-        tflite_results = [np.array([[-113, 113]], dtype=np.int8)]
+        tflite_results = tflite_preprocessor.get_raw_model_output()
         compare_results(mlir_results, tflite_results, output_details)
         # print(mlir_results)
 
 
-class PersonDetectionTfliteModuleTest(unittest.TestCase):
+class AlbertTfliteModuleTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def configure(self, pytestconfig):
         self.save_mlir = pytestconfig.getoption("save_mlir")
         self.save_vmfb = pytestconfig.getoption("save_vmfb")
 
     def setUp(self):
-        self.module_tester = PersonDetectionTfliteModuleTester(self)
+        self.module_tester = AlbertTfliteModuleTester(self)
         self.module_tester.save_mlir = self.save_mlir
 
-    @pytest.mark.skip(reason="TFLite is broken with this model")
+    import sys
+
+    @pytest.mark.xfail(
+        sys.platform == "darwin", reason="known macos tflite install issue"
+    )
     def test_module_static_cpu(self):
         self.module_tester.dynamic = False
         self.module_tester.device = "cpu"
         self.module_tester.create_and_check_module()
 
 
 if __name__ == "__main__":
-    # module_tester = PersonDetectionTfliteModuleTester()
+    # module_tester = AlbertTfliteModuleTester()
     # module_tester.save_mlir = True
     # module_tester.save_vmfb = True
     # module_tester.create_and_check_module()

tank/albert_lite_base/albert_lite_base_tflite_mlir_test.py

@@ -1,5 +1,5 @@
 import numpy as np
-from shark.shark_importer import SharkImporter
+from shark.shark_downloader import SharkDownloader
 from shark.shark_inference import SharkInference
 import pytest
 import unittest
@@ -70,58 +70,26 @@ def __init__(
     def create_and_check_module(self):
         shark_args.save_mlir = self.save_mlir
         shark_args.save_vmfb = self.save_vmfb
-
-        # Preprocess to get SharkImporter input args
-        tflite_preprocessor = TFLitePreprocessor(model_name="albert_lite_base")
-        raw_model_file_path = tflite_preprocessor.get_raw_model_file()
-        inputs = tflite_preprocessor.get_inputs()
-        tflite_interpreter = tflite_preprocessor.get_interpreter()
-
-        # Use SharkImporter to get SharkInference input args
-        my_shark_importer = SharkImporter(
-            module=tflite_interpreter,
-            inputs=inputs,
-            frontend="tflite",
-            raw_model_file=raw_model_file_path,
-        )
-        mlir_model, func_name = my_shark_importer.import_mlir()
-
-        # Use SharkInference to get inference result
-        shark_module = SharkInference(
-            mlir_module=mlir_model,
-            function_name=func_name,
-            device=self.device,
-            mlir_dialect="tflite",
+        shark_downloader = SharkDownloader(
+            model_name="albert_lite_base",
+            tank_url="https://storage.googleapis.com/shark_tank",
+            local_tank_dir="./../gen_shark_tank",
+            model_type="tflite",
+            input_json="input.json",
+            input_type="int32",
         )
-
-        # Case1: Use shark_importer default generate inputs
-        shark_module.compile()
-        mlir_results = shark_module.forward(inputs)
-        ## post process results for compare
-        input_details, output_details = tflite_preprocessor.get_model_details()
-        mlir_results = list(mlir_results)
-        for i in range(len(output_details)):
-            dtype = output_details[i]["dtype"]
-            mlir_results[i] = mlir_results[i].astype(dtype)
-        tflite_results = tflite_preprocessor.get_raw_model_output()
-        compare_results(mlir_results, tflite_results, output_details)
-
-        # Case2: Use manually set inputs
-        input_details, output_details = tflite_preprocessor.get_model_details()
-        inputs = generate_inputs(input_details)  # new inputs
+        tflite_tosa_model = shark_downloader.get_mlir_file()
+        inputs = shark_downloader.get_inputs()
 
         shark_module = SharkInference(
-            mlir_module=mlir_model,
-            function_name=func_name,
+            mlir_module=tflite_tosa_model,
+            function_name="main",
             device=self.device,
             mlir_dialect="tflite",
         )
         shark_module.compile()
-        mlir_results = shark_module.forward(inputs)
-        ## post process results for compare
-        tflite_results = tflite_preprocessor.get_raw_model_output()
-        compare_results(mlir_results, tflite_results, output_details)
-        # print(mlir_results)
+        shark_module.forward(inputs)
+        # print(shark_results)
 
 
 class AlbertTfliteModuleTest(unittest.TestCase):
@@ -146,9 +114,24 @@ def test_module_static_cpu(self):
 
 
 if __name__ == "__main__":
+    unittest.main()
     # module_tester = AlbertTfliteModuleTester()
-    # module_tester.save_mlir = True
-    # module_tester.save_vmfb = True
     # module_tester.create_and_check_module()
 
-    unittest.main()
+# TEST RESULT:
+# (shark.venv) nod% python albert_lite_base_tflite_mlir_test.py
+# load json inputs
+# TMP_MODEL_DIR = shark/SHARK/shark/./../gen_shark_tank/tflite
+# Model has not been download.shark_downloader will automatically download by tank_url if provided. You can also manually to download the model from shark_tank by yourself.
+# TMP_MODELNAME_DIR = shark/SHARK/shark/./../gen_shark_tank/tflite/albert_lite_base
+# Download mlir model https://storage.googleapis.com/shark_tank/tflite/albert_lite_base/albert_lite_base_tosa.mlir
+# Get tosa.mlir model return
+# Target triple found:x86_64-linux-gnu
+# (shark.venv) nod% python albert_lite_base_tflite_mlir_test.py
+# load json inputs
+# TMP_MODEL_DIR = shark/SHARK/shark/./../gen_shark_tank/tflite
+# TMP_MODELNAME_DIR = shark/SHARK/shark/./../gen_shark_tank/tflite/albert_lite_base
+# Model has been downloaded before. shark/SHARK/shark/./../gen_shark_tank/tflite/albert_lite_base/albert_lite_base_tosa.mlir
+# Get tosa.mlir model return
+# Target triple found:x86_64-linux-gnu
+#