Remove tf package dep for SharkDownloader tflite tests (huggingface#212)

Author: AmosLewis

generate_sharktank.py

@@ -115,6 +115,7 @@ def save_tflite_model(tflite_model_list):
     with open(tflite_model_list) as csvfile:
         tflite_reader = csv.reader(csvfile, delimiter=",")
         for row in tflite_reader:
+            print("\n")
             tflite_model_name = row[0]
             tflite_model_link = row[1]
             print("tflite_model_name", tflite_model_name)
@@ -125,13 +126,6 @@ def save_tflite_model(tflite_model_list):
             os.makedirs(tflite_model_name_dir, exist_ok=True)
             print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")
 
-            tflite_tosa_file = "/".join(
-                [
-                    tflite_model_name_dir,
-                    str(tflite_model_name) + "_tflite.mlir",
-                ]
-            )
-
             # Preprocess to get SharkImporter input args
             tflite_preprocessor = TFLitePreprocessor(str(tflite_model_name))
             raw_model_file_path = tflite_preprocessor.get_raw_model_file()
@@ -145,15 +139,11 @@ def save_tflite_model(tflite_model_list):
                 frontend="tflite",
                 raw_model_file=raw_model_file_path,
             )
-            mlir_model, func_name = my_shark_importer.import_mlir()
-
-            if os.path.exists(tflite_tosa_file):
-                print("Exists", tflite_tosa_file)
-            else:
-                mlir_str = mlir_model.decode("utf-8")
-                with open(tflite_tosa_file, "w") as f:
-                    f.write(mlir_str)
-                print(f"Saved mlir in {tflite_tosa_file}")
+            my_shark_importer.import_debug(
+                dir=tflite_model_name_dir,
+                model_name=tflite_model_name,
+                func_name="main",
+            )
 
 
 # Validates whether the file is present or not.
@@ -170,7 +160,7 @@ def is_valid_file(arg):
         "--torch_model_csv",
         type=lambda x: is_valid_file(x),
         default="./tank/pytorch/torch_model_list.csv",
-        help="""Contains the file with torch_model name and args. 
+        help="""Contains the file with torch_model name and args.
              Please see: https://github.com/nod-ai/SHARK/blob/main/tank/pytorch/torch_model_list.csv""",
     )
     parser.add_argument(

shark/shark_downloader.py

@@ -27,7 +27,7 @@
     "int8": np.int8,
 }
 
-WORKDIR = os.path.join(os.path.dirname(__file__), "gen_shark_tank")
+WORKDIR = os.path.join(os.path.dirname(__file__), "./../gen_shark_tank")
 
 # Checks whether the directory and files exists.
 def check_dir_exists(model_name, frontend="torch", dynamic=""):
@@ -83,6 +83,36 @@ def download_torch_model(model_name, dynamic=False):
     return mlir_file, function_name, inputs_tuple, golden_out_tuple
 
 
+# Downloads the tflite model from gs://shark_tank dir.
+def download_tflite_model(model_name, dynamic=False):
+    dyn_str = "_dynamic" if dynamic else ""
+    os.makedirs(WORKDIR, exist_ok=True)
+    if not check_dir_exists(model_name, dyn_str):
+        gs_command = (
+            'gsutil -o "GSUtil:parallel_process_count=1" cp -r gs://shark_tank'
+            + "/"
+            + model_name
+            + " "
+            + WORKDIR
+        )
+        if os.system(gs_command) != 0:
+            raise Exception("model not present in the tank. Contact Nod Admin")
+
+    model_dir = os.path.join(WORKDIR, model_name)
+    with open(
+        os.path.join(model_dir, model_name + dyn_str + "_tflite.mlir")
+    ) as f:
+        mlir_file = f.read()
+
+    function_name = str(np.load(os.path.join(model_dir, "function_name.npy")))
+    inputs = np.load(os.path.join(model_dir, "inputs.npz"))
+    golden_out = np.load(os.path.join(model_dir, "golden_out.npz"))
+
+    inputs_tuple = tuple([inputs[key] for key in inputs])
+    golden_out_tuple = tuple([golden_out[key] for key in golden_out])
+    return mlir_file, function_name, inputs_tuple, golden_out_tuple
+
+
 def download_tf_model(model_name):
     model_name = model_name.replace("/", "_")
     os.makedirs(WORKDIR, exist_ok=True)
@@ -109,149 +139,3 @@ def download_tf_model(model_name):
     inputs_tuple = tuple([inputs[key] for key in inputs])
     golden_out_tuple = tuple([golden_out[key] for key in golden_out])
     return mlir_file, function_name, inputs_tuple, golden_out_tuple
-
-
-class SharkDownloader:
-    def __init__(
-        self,
-        model_name: str,
-        tank_url: str = "https://storage.googleapis.com/shark_tank",
-        local_tank_dir: str = "./../gen_shark_tank/tflite",
-        model_type: str = "tflite",
-        input_json: str = "input.json",
-        input_type: str = "int32",
-    ):
-        self.model_name = model_name
-        self.local_tank_dir = local_tank_dir
-        self.tank_url = tank_url
-        self.model_type = model_type
-        self.input_json = input_json  # optional if you don't have input
-        self.input_type = input_type_to_np_dtype[
-            input_type
-        ]  # optional if you don't have input
-        self.mlir_file = None  # .mlir file local address.
-        self.mlir_url = None
-        self.inputs = None  # Input has to be (list of np.array) for sharkInference.forward use
-        self.mlir_model = []
-
-        # create tmp model file directory
-        if self.tank_url is None and self.model_name is None:
-            print("Error. No tank_url, No model name,Please input either one.")
-            return
-
-        self.workdir = os.path.join(
-            os.path.dirname(__file__), self.local_tank_dir
-        )
-        os.makedirs(self.workdir, exist_ok=True)
-        print(f"TMP_MODEL_DIR = {self.workdir}")
-        # use model name get dir.
-        self.model_name_dir = os.path.join(self.workdir, str(self.model_name))
-        if not os.path.exists(self.model_name_dir):
-            print(
-                "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."
-            )
-        os.makedirs(self.model_name_dir, exist_ok=True)
-        print(f"TMP_MODELNAME_DIR = {self.model_name_dir}")
-
-        # read inputs from json file
-        self.load_json_input()
-        # get milr model file
-        self.load_mlir_model()
-
-    def get_mlir_file(self):
-        return self.mlir_model
-
-    def get_inputs(self):
-        return self.inputs
-
-    def load_json_input(self):
-        print("load json inputs")
-        if self.model_type in ["tflite"]:
-            input_url = (
-                self.tank_url + "/" + str(self.model_name) + "/" + "input.json"
-            )
-            input_file = "/".join([self.model_name_dir, str(self.input_json)])
-            if os.path.exists(input_file):
-                print("Input has been downloaded before.", input_file)
-            else:
-                print("Download input", input_url)
-                urllib.request.urlretrieve(input_url, input_file)
-
-            args = []
-            with open(input_file, "r") as f:
-                args = json.load(f)
-            self.inputs = [
-                np.asarray(arg, dtype=self.input_type) for arg in args
-            ]
-        else:
-            print(
-                "No json input required for current model type. "
-                "You could call setup_inputs(YOU_INPUTS)."
-            )
-        return self.inputs
-
-    def load_mlir_model(self):
-        if self.model_type in ["tflite"]:
-            self.mlir_url = (
-                self.tank_url
-                + "/"
-                + str(self.model_name)
-                + "/"
-                + str(self.model_name)
-                + "_tflite.mlir"
-            )
-            self.mlir_file = "/".join(
-                [self.model_name_dir, str(self.model_name) + "_tfite.mlir"]
-            )
-        elif self.model_type in ["tensorflow"]:
-            self.mlir_url = (
-                self.tank_url
-                + "/"
-                + str(self.model_name)
-                + "/"
-                + str(self.model_name)
-                + "_tf.mlir"
-            )
-            self.mlir_file = "/".join(
-                [self.model_name_dir, str(self.model_name) + "_tf.mlir"]
-            )
-        elif self.model_type in ["torch", "jax", "mhlo", "tosa"]:
-            self.mlir_url = (
-                self.tank_url
-                + "/"
-                + str(self.model_name)
-                + "/"
-                + str(self.model_name)
-                + "_"
-                + str(self.model_type)
-                + ".mlir"
-            )
-            self.mlir_file = "/".join(
-                [
-                    self.model_name_dir,
-                    str(self.model_name)
-                    + "_"
-                    + str(self.model_type)
-                    + ".mlir",
-                ]
-            )
-        else:
-            print("Unsupported mlir model")
-
-        if os.path.exists(self.mlir_file):
-            print("Model has been downloaded before.", self.mlir_file)
-        else:
-            print("Download mlir model", self.mlir_url)
-            urllib.request.urlretrieve(self.mlir_url, self.mlir_file)
-
-        print("Get .mlir model return")
-        with open(self.mlir_file) as f:
-            self.mlir_model = f.read()
-        return self.mlir_model
-
-    def setup_inputs(self, inputs):
-        print("Setting up inputs. Input has to be (list of np.array)")
-        self.inputs = inputs

shark/shark_importer.py

@@ -129,7 +129,7 @@ def save_data(
         inputs_name = "inputs.npz"
         outputs_name = "golden_out.npz"
         func_file_name = "function_name"
-        model_name_mlir = model_name + ".mlir"
+        model_name_mlir = model_name + "_" + self.frontend + ".mlir"
         np.savez(os.path.join(dir, inputs_name), *inputs)
         np.savez(os.path.join(dir, outputs_name), *outputs)
         np.save(os.path.join(dir, func_file_name), np.array(func_name))
@@ -139,6 +139,8 @@ def save_data(
             mlir_str = mlir_data.operation.get_asm()
         elif self.frontend == "tf":
             mlir_str = mlir_data.decode("utf-8")
+        elif self.frontend == "tflite":
+            mlir_str = mlir_data.decode("utf-8")
         with open(os.path.join(dir, model_name_mlir), "w") as mlir_file:
             mlir_file.write(mlir_str)
 
@@ -214,6 +216,14 @@ def import_debug(
         if self.frontend in ["tflite", "tf-lite"]:
             # TODO(Chi): Validate it for tflite models.
             golden_out = self.module.invoke_tflite(self.inputs)
+            self.save_data(
+                dir,
+                model_name,
+                imported_mlir[0],
+                imported_mlir[1],
+                self.inputs,
+                golden_out,
+            )
             return (
                 imported_mlir,
                 self.inputs,

shark/tests/test_shark_importer.py

@@ -98,7 +98,7 @@ def create_and_check_module(self):
         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()
+        tflite_results = tflite_preprocessor.get_golden_output()
         compare_results(mlir_results, tflite_results, output_details)
 
         # Case2: Use manually set inputs
@@ -114,7 +114,7 @@ def create_and_check_module(self):
         shark_module.compile()
         mlir_results = shark_module.forward(inputs)
         ## post process results for compare
-        tflite_results = tflite_preprocessor.get_raw_model_output()
+        tflite_results = tflite_preprocessor.get_golden_output()
         compare_results(mlir_results, tflite_results, output_details)
         # print(mlir_results)
 

shark/tflite_utils.py

@@ -3,7 +3,6 @@
 import os
 import csv
 import urllib.request
-import json
 
 
 class TFLiteModelUtil:
@@ -48,6 +47,8 @@ def invoke_tflite(self, inputs):
             )
 
         for i in range(len(self.output_details)):
+            # print("output_details ", i, "shape", self.output_details[i]["shape"].__name__,
+            #       ", dtype: ", self.output_details[i]["dtype"].__name__)
             out_dtype = self.output_details[i]["dtype"]
             tflite_results[i] = tflite_results[i].astype(out_dtype)
         return tflite_results
@@ -84,6 +85,7 @@ def __init__(
             None  # could be tflite/tf/torch_interpreter in utils
         )
         self.input_file = None
+        self.output_file = None
 
         # create tmp model file directory
         if self.model_path is None and self.model_name is None:
@@ -127,7 +129,9 @@ def load_tflite_model(self):
         self.mlir_file = "/".join(
             [tflite_model_name_dir, str(self.model_name) + "_tflite.mlir"]
         )
-        self.input_file = "/".join([tflite_model_name_dir, "input.json"])
+        self.input_file = "/".join([tflite_model_name_dir, "inputs"])
+        self.output_file = "/".join([tflite_model_name_dir, "golden_out"])
+        # np.save("/".join([tflite_model_name_dir, "function_name"]), np.array("main"))
 
         if os.path.exists(self.raw_model_file):
             print(
@@ -165,21 +169,15 @@ def setup_interpreter(self):
     def generate_inputs(self, input_details):
         self.inputs = []
         for tmp_input in input_details:
-            # print(str(tmp_input["shape"]), tmp_input["dtype"].__name__)
+            print(
+                "input_details shape:",
+                str(tmp_input["shape"]),
+                " type:",
+                tmp_input["dtype"].__name__,
+            )
             self.inputs.append(
                 np.ones(shape=tmp_input["shape"], dtype=tmp_input["dtype"])
             )
-        # save inputs into json file
-        tmp_json = []
-        for tmp_input in input_details:
-            # print(str(tmp_input["shape"]), tmp_input["dtype"].__name__)
-            tmp_json.append(
-                np.ones(
-                    shape=tmp_input["shape"], dtype=tmp_input["dtype"]
-                ).tolist()
-            )
-        with open(self.input_file, "w") as f:
-            json.dump(tmp_json, f)
         return self.inputs
 
     def setup_inputs(self, inputs):
@@ -195,8 +193,9 @@ def get_mlir_file(self):
     def get_inputs(self):
         return self.inputs
 
-    def get_raw_model_output(self):
+    def get_golden_output(self):
         self.output_tensor = self.interpreter.invoke_tflite(self.inputs)
+        np.savez(self.output_file, *self.output_tensor)
         return self.output_tensor
 
     def get_model_details(self):