diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index cc0570a3..7b63f3be 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -53,7 +53,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install flake8 pytest yapf toml + python -m pip install flake8 pytest toml if [ -f requirements.txt ]; then pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://github.com/llvm/torch-mlir/releases -f https://github.com/nod-ai/SHARK-Runtime/releases; fi - name: Lint with flake8 run: | @@ -61,7 +61,6 @@ jobs: flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude shark.venv,lit.cfg.py # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude shark.venv,lit.cfg.py - yapf -i --style .style.yapf shark/*.py - name: Build and validate the package run: | diff --git a/.github/workflows/python-format.yml b/.github/workflows/python-format.yml new file mode 100644 index 00000000..e66de94a --- /dev/null +++ b/.github/workflows/python-format.yml @@ -0,0 +1,15 @@ +name: black-formatter +on: [pull_request] +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v1 + - name: Set up Python 3.10 + uses: actions/setup-python@v1 + with: + python-version: 3.10.5 + - name: Install Black + run: pip install black + - name: Run formatter check on the entire project. + run: black --line-length 80 --check . diff --git a/.github/workflows/test-models.yml b/.github/workflows/test-models.yml index 6659e8f3..94f0a714 100644 --- a/.github/workflows/test-models.yml +++ b/.github/workflows/test-models.yml @@ -37,7 +37,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install flake8 pytest yapf toml + python -m pip install flake8 pytest toml - name: Lint with flake8 run: | @@ -45,7 +45,6 @@ jobs: flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude lit.cfg.py # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude lit.cfg.py - yapf -i --style .style.yapf shark/*.py - name: Validate Models run: | diff --git a/benchmarks/hf_model_benchmark.py b/benchmarks/hf_model_benchmark.py index 000497a7..e8526286 100644 --- a/benchmarks/hf_model_benchmark.py +++ b/benchmarks/hf_model_benchmark.py @@ -6,16 +6,16 @@ parser.add_argument( "--model_name", type=str, required=True, - help= - "Specifies name of HF model to benchmark. (For exmaple \"microsoft/MiniLM-L12-H384-uncased\"" + help='Specifies name of HF model to benchmark. (For exmaple "microsoft/MiniLM-L12-H384-uncased"', ) load_args, unknown = parser.parse_known_args() if __name__ == "__main__": model_name = load_args.model_name test_input = torch.randint(2, (1, 128)) - shark_module = SharkHFBenchmarkRunner(model_name, (test_input,), - jit_trace=True) + shark_module = SharkHFBenchmarkRunner( + model_name, (test_input,), jit_trace=True + ) shark_module.benchmark_c() shark_module.benchmark_python((test_input,)) shark_module.benchmark_torch(test_input) diff --git a/benchmarks/hf_transformer.py b/benchmarks/hf_transformer.py index fd4976d5..5f4ae9d9 100644 --- a/benchmarks/hf_transformer.py +++ b/benchmarks/hf_transformer.py @@ -1,8 +1,12 @@ import torch -from shark.shark_runner import SharkBenchmarkRunner +from shark.shark_benchmark_runner import SharkBenchmarkRunner from shark.parser import shark_args from transformers import AutoTokenizer, AutoModelForSequenceClassification -from onnxruntime.transformers.benchmark import run_pytorch, run_tensorflow, run_onnxruntime +from onnxruntime.transformers.benchmark import ( + run_pytorch, + run_tensorflow, + run_onnxruntime, +) from onnxruntime.transformers.huggingface_models import MODELS from onnxruntime.transformers.benchmark_helper import ConfigModifier, Precision import os @@ -10,7 +14,6 @@ import psutil class OnnxFusionOptions(object): - def __init__(self): self.disable_gelu = False self.disable_layer_norm = False @@ -25,17 +28,13 @@ class OnnxFusionOptions(object): class HuggingFaceLanguage(torch.nn.Module): - def __init__(self, hf_model_name): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( hf_model_name, # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -62,8 +61,16 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner): ) self.model_name = model_name model = HuggingFaceLanguage(model_name) - SharkBenchmarkRunner.__init__(self, model, input, dynamic, self.device, - jit_trace, from_aot, frontend) + SharkBenchmarkRunner.__init__( + self, + model, + input, + dynamic, + self.device, + jit_trace, + from_aot, + frontend, + ) def benchmark_torch(self, inputs): use_gpu = self.device == "gpu" @@ -74,10 +81,20 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner): sequence_lengths = [inputs.shape[-1]] cache_dir = os.path.join(".", "cache_models") verbose = False - result = run_pytorch(use_gpu, [self.model_name], None, config_modifier, - Precision.FLOAT32, num_threads, batch_sizes, - sequence_lengths, shark_args.num_iterations, False, - cache_dir, verbose) + result = run_pytorch( + use_gpu, + [self.model_name], + None, + config_modifier, + Precision.FLOAT32, + num_threads, + batch_sizes, + sequence_lengths, + shark_args.num_iterations, + False, + cache_dir, + verbose, + ) print( f"ONNX Pytorch-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}" ) @@ -92,10 +109,19 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner): sequence_lengths = [inputs.shape[-1]] cache_dir = os.path.join(".", "cache_models") verbose = False - result = run_tensorflow(use_gpu, [self.model_name], None, - config_modifier, Precision.FLOAT32, num_threads, - batch_sizes, sequence_lengths, - shark_args.num_iterations, cache_dir, verbose) + result = run_tensorflow( + use_gpu, + [self.model_name], + None, + config_modifier, + Precision.FLOAT32, + num_threads, + batch_sizes, + sequence_lengths, + shark_args.num_iterations, + cache_dir, + verbose, + ) print( f"ONNX TF-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}" ) @@ -105,7 +131,8 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner): print( f"{self.model_name} is currently not supported in ORT's HF. Check \ https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/huggingface_models.py \ -for currently supported models. Exiting benchmark ONNX.") +for currently supported models. Exiting benchmark ONNX." + ) return use_gpu = self.device == "gpu" num_threads = psutil.cpu_count(logical=False) @@ -121,17 +148,34 @@ for currently supported models. Exiting benchmark ONNX.") use_raw_attention_mask = True model_fusion_statistics = {} overwrite = False - model_source = "pt" #Either "pt" or "tf" + model_source = "pt" # Either "pt" or "tf" provider = None config_modifier = ConfigModifier(None) onnx_args = OnnxFusionOptions() result = run_onnxruntime( - use_gpu, provider, [self.model_name], None, config_modifier, - Precision.FLOAT32, num_threads, batch_sizes, sequence_lengths, - shark_args.num_iterations, input_counts, optimize_onnx, - validate_onnx, cache_dir, onnx_dir, verbose, overwrite, - disable_ort_io_binding, use_raw_attention_mask, - model_fusion_statistics, model_source, onnx_args) + use_gpu, + provider, + [self.model_name], + None, + config_modifier, + Precision.FLOAT32, + num_threads, + batch_sizes, + sequence_lengths, + shark_args.num_iterations, + input_counts, + optimize_onnx, + validate_onnx, + cache_dir, + onnx_dir, + verbose, + overwrite, + disable_ort_io_binding, + use_raw_attention_mask, + model_fusion_statistics, + model_source, + onnx_args, + ) print( f"ONNX ORT-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}" ) diff --git a/benchmarks/tests/test_benchmark.py b/benchmarks/tests/test_benchmark.py index 39adc656..0bc9a214 100644 --- a/benchmarks/tests/test_benchmark.py +++ b/benchmarks/tests/test_benchmark.py @@ -1,19 +1,23 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers import torch import tensorflow as tf import numpy as np import torchvision.models as models -from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel +from transformers import ( + AutoModelForSequenceClassification, + BertTokenizer, + TFBertModel, +) import importlib import pytest import unittest torch.manual_seed(0) -gpus = tf.config.experimental.list_physical_devices('GPU') +gpus = tf.config.experimental.list_physical_devices("GPU") for gpu in gpus: - tf.config.experimental.set_memory_growth(gpu, True) + tf.config.experimental.set_memory_growth(gpu, True) ##################### Tensorflow Hugging Face LM Models ################################### MAX_SEQUENCE_LENGTH = 512 @@ -23,12 +27,11 @@ BATCH_SIZE = 1 tf_bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] class TFHuggingFaceLanguage(tf.Module): - def __init__(self, hf_model_name): super(TFHuggingFaceLanguage, self).__init__() # Create a BERT trainer with the created network. @@ -36,7 +39,8 @@ class TFHuggingFaceLanguage(tf.Module): # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=tf_bert_input) def forward(self, input_ids, attention_mask, token_type_ids): @@ -47,15 +51,21 @@ def get_TFhf_model(name): model = TFHuggingFaceLanguage(name) tokenizer = BertTokenizer.from_pretrained(name) text = "Replace me by any text you'd like." - encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + encoded_input = tokenizer( + text, + padding="max_length", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) - test_input = (encoded_input["input_ids"], encoded_input["attention_mask"], - encoded_input["token_type_ids"]) + tf.convert_to_tensor(encoded_input[key]), 0 + ) + test_input = ( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ) actual_out = model.forward(*test_input) return model, test_input, actual_out @@ -64,17 +74,13 @@ def get_TFhf_model(name): class HuggingFaceLanguage(torch.nn.Module): - def __init__(self, hf_model_name): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( hf_model_name, # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -96,7 +102,6 @@ def get_hf_model(name): class VisionModule(torch.nn.Module): - def __init__(self, model): super().__init__() self.model = model @@ -117,46 +122,56 @@ def get_vision_model(torch_model): ############################# Benchmark Tests #################################### pytest_benchmark_param = pytest.mark.parametrize( - ('dynamic', 'device'), + ("dynamic", "device"), [ - pytest.param(False, 'cpu'), + pytest.param(False, "cpu"), # TODO: Language models are failing for dynamic case.. - pytest.param(True, 'cpu', marks=pytest.mark.skip), - pytest.param(False, - 'gpu', - marks=pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found")), - pytest.param(True, - 'gpu', - marks=pytest.mark.skip), + pytest.param(True, "cpu", marks=pytest.mark.skip), pytest.param( False, - 'vulkan', + "gpu", + marks=pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ), + ), + pytest.param(True, "gpu", marks=pytest.mark.skip), + pytest.param( + False, + "vulkan", marks=pytest.mark.skipif( check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - )), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ), + ), pytest.param( True, - 'vulkan', + "vulkan", marks=pytest.mark.skipif( check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - )), - ]) + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ), + ), + ], +) -@pytest.mark.skipif(importlib.util.find_spec("iree.tools") is None, - reason="Cannot find tools to import TF") +@pytest.mark.skipif( + importlib.util.find_spec("iree.tools") is None, + reason="Cannot find tools to import TF", +) @pytest_benchmark_param def test_bench_minilm_torch(dynamic, device): model, test_input, act_out = get_hf_model( - "microsoft/MiniLM-L12-H384-uncased") - shark_module = SharkInference(model, (test_input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=True) + "microsoft/MiniLM-L12-H384-uncased" + ) + shark_module = SharkInference( + model, + (test_input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=True, + ) try: # If becnhmarking succesful, assert success/True. shark_module.compile() @@ -167,17 +182,21 @@ def test_bench_minilm_torch(dynamic, device): assert False -@pytest.mark.skipif(importlib.util.find_spec("iree.tools") is None, - reason="Cannot find tools to import TF") +@pytest.mark.skipif( + importlib.util.find_spec("iree.tools") is None, + reason="Cannot find tools to import TF", +) @pytest_benchmark_param def test_bench_distilbert(dynamic, device): model, test_input, act_out = get_TFhf_model("distilbert-base-uncased") - shark_module = SharkInference(model, - test_input, - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=True) + shark_module = SharkInference( + model, + test_input, + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=True, + ) try: # If becnhmarking succesful, assert success/True. shark_module.set_frontend("tensorflow") @@ -193,12 +212,14 @@ def test_bench_distilbert(dynamic, device): @pytest_benchmark_param def test_bench_xlm_roberta(dynamic, device): model, test_input, act_out = get_TFhf_model("xlm-roberta-base") - shark_module = SharkInference(model, - test_input, - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=True) + shark_module = SharkInference( + model, + test_input, + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=True, + ) try: # If becnhmarking succesful, assert success/True. shark_module.set_frontend("tensorflow") diff --git a/benchmarks/tests/test_hf_benchmark.py b/benchmarks/tests/test_hf_benchmark.py index a19a0a5a..cfbd6d38 100644 --- a/benchmarks/tests/test_hf_benchmark.py +++ b/benchmarks/tests/test_hf_benchmark.py @@ -9,25 +9,31 @@ torch.manual_seed(0) # Test running benchmark module without failing. pytest_benchmark_param = pytest.mark.parametrize( - ('dynamic', 'device'), + ("dynamic", "device"), [ - pytest.param(False, 'cpu'), + pytest.param(False, "cpu"), # TODO: Language models are failing for dynamic case.. - pytest.param(True, 'cpu', marks=pytest.mark.skip), - ]) + pytest.param(True, "cpu", marks=pytest.mark.skip), + ], +) -@pytest.mark.skipif(importlib.util.find_spec("onnxruntime") is None, - reason="Cannot find ONNXRUNTIME.") +@pytest.mark.skipif( + importlib.util.find_spec("onnxruntime") is None, + reason="Cannot find ONNXRUNTIME.", +) @pytest_benchmark_param def test_HFbench_minilm_torch(dynamic, device): model_name = "bert-base-uncased" test_input = torch.randint(2, (1, 128)) try: - shark_module = SharkHFBenchmarkRunner(model_name, (test_input,), - jit_trace=True, - dynamic=dynamic, - device=device) + shark_module = SharkHFBenchmarkRunner( + model_name, + (test_input,), + jit_trace=True, + dynamic=dynamic, + device=device, + ) shark_module.benchmark_c() shark_module.benchmark_python((test_input,)) shark_module.benchmark_torch(test_input) diff --git a/generate_sharktank.py b/generate_sharktank.py index 5966d8ec..d3d274a5 100644 --- a/generate_sharktank.py +++ b/generate_sharktank.py @@ -18,13 +18,15 @@ import argparse import iree.compiler.tflite as ireec_tflite from shark.iree_utils import IREE_TARGET_MAP -class SharkTank: - def __init__(self, - torch_model_list: str = None, - tf_model_list: str = None, - tflite_model_list: str = None, - upload: bool = False): +class SharkTank: + def __init__( + self, + torch_model_list: str = None, + tf_model_list: str = None, + tflite_model_list: str = None, + upload: bool = False, + ): self.torch_model_list = torch_model_list self.tf_model_list = tf_model_list self.tflite_model_list = tflite_model_list @@ -43,30 +45,53 @@ class SharkTank: # compile and run tfhub tflite if self.tflite_model_list is not None: print("Setting up for tflite TMP_DIR") - self.tflite_workdir = os.path.join(os.path.dirname(__file__), "./gen_shark_tank/tflite") + self.tflite_workdir = os.path.join( + os.path.dirname(__file__), "./gen_shark_tank/tflite" + ) print(f"tflite TMP_shark_tank_DIR = {self.tflite_workdir}") os.makedirs(self.tflite_workdir, exist_ok=True) with open(self.tflite_model_list) as csvfile: - tflite_reader = csv.reader(csvfile, delimiter=',') + tflite_reader = csv.reader(csvfile, delimiter=",") for row in tflite_reader: tflite_model_name = row[0] tflite_model_link = row[1] print("tflite_model_name", tflite_model_name) print("tflite_model_link", tflite_model_link) - tflite_model_name_dir = os.path.join(self.tflite_workdir, str(tflite_model_name)) + tflite_model_name_dir = os.path.join( + self.tflite_workdir, str(tflite_model_name) + ) os.makedirs(tflite_model_name_dir, exist_ok=True) - tflite_saving_file = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_tflite.tflite']) - tflite_tosa_file = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_tosa.mlir']) - self.binary = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_module.bytecode']) - print("Setting up local address for tflite model file: ", tflite_saving_file) + tflite_saving_file = "/".join( + [ + tflite_model_name_dir, + str(tflite_model_name) + "_tflite.tflite", + ] + ) + tflite_tosa_file = "/".join( + [ + tflite_model_name_dir, + str(tflite_model_name) + "_tosa.mlir", + ] + ) + self.binary = "/".join( + [ + tflite_model_name_dir, + str(tflite_model_name) + "_module.bytecode", + ] + ) + print( + "Setting up local address for tflite model file: ", + tflite_saving_file, + ) if os.path.exists(tflite_saving_file): print(tflite_saving_file, "exists") else: print("Download tflite model") - urllib.request.urlretrieve(str(tflite_model_link), - tflite_saving_file) + urllib.request.urlretrieve( + str(tflite_model_link), tflite_saving_file + ) if os.path.exists(tflite_tosa_file): print("Exists", tflite_tosa_file) @@ -76,19 +101,35 @@ class SharkTank: tflite_saving_file, input_type="tosa", save_temp_iree_input=tflite_tosa_file, - target_backends=[IREE_TARGET_MAP['cpu']], - import_only=False) + target_backends=[IREE_TARGET_MAP["cpu"]], + import_only=False, + ) if self.upload == True: print("upload tmp tank to gcp") - os.system('gsutil cp -r ./gen_shark_tank gs://shark_tank/') -if __name__ == '__main__': + os.system("gsutil cp -r ./gen_shark_tank gs://shark_tank/") + + +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--torch_model_list", type=str, default="./tank/torch/torch_model_list.csv") - parser.add_argument("--tf_model_list", type=str, default="./tank/tf/tf_model_list.csv") - parser.add_argument("--tflite_model_list", type=str, default="./tank/tflite/tflite_model_list.csv") + parser.add_argument( + "--torch_model_list", + type=str, + default="./tank/torch/torch_model_list.csv", + ) + parser.add_argument( + "--tf_model_list", type=str, default="./tank/tf/tf_model_list.csv" + ) + parser.add_argument( + "--tflite_model_list", + type=str, + default="./tank/tflite/tflite_model_list.csv", + ) parser.add_argument("--upload", type=bool, default=False) args = parser.parse_args() - SharkTank(torch_model_list=args.torch_model_list, tf_model_list=args.tf_model_list, - tflite_model_list=args.tflite_model_list, upload=args.upload) - + SharkTank( + torch_model_list=args.torch_model_list, + tf_model_list=args.tf_model_list, + tflite_model_list=args.tflite_model_list, + upload=args.upload, + ) diff --git a/setup.py b/setup.py index 96fc890a..6a1a118b 100644 --- a/setup.py +++ b/setup.py @@ -26,7 +26,7 @@ setup( "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ], - packages=find_packages(exclude=('examples')), + packages=find_packages(exclude=("examples")), python_requires=">=3.7", install_requires=[ "numpy", diff --git a/shark/backward_makefx.py b/shark/backward_makefx.py index 79c0d013..8d1b55eb 100644 --- a/shark/backward_makefx.py +++ b/shark/backward_makefx.py @@ -22,7 +22,6 @@ import tempfile class MakeFxModule: - def __init__(self, model, inputs, labels=None, custom_inference_fn=None): self.model = model self.inputs = inputs @@ -52,20 +51,28 @@ class MakeFxModule: return fx_g def generate_graph(self): - fx_g = make_fx(self.custom_inference_fn, - decomposition_table=get_decompositions([ - torch.ops.aten.embedding_dense_backward, - torch.ops.aten.native_layer_norm_backward, - torch.ops.aten.slice_backward, - torch.ops.aten.select_backward - ]))(dict(self.model.named_parameters()), - dict(self.model.named_buffers()), self.inputs) + fx_g = make_fx( + self.custom_inference_fn, + decomposition_table=get_decompositions( + [ + torch.ops.aten.embedding_dense_backward, + torch.ops.aten.native_layer_norm_backward, + torch.ops.aten.slice_backward, + torch.ops.aten.select_backward, + ] + ), + )( + dict(self.model.named_parameters()), + dict(self.model.named_buffers()), + self.inputs, + ) fx_g.graph.set_codegen(torch.fx.graph.CodeGen()) fx_g.recompile() fx_g = self.change_fx_graph_return_to_tuple(fx_g) ts_g = torch.jit.script(fx_g) - temp = tempfile.NamedTemporaryFile(suffix='_shark_ts', - prefix='temp_ts_') + temp = tempfile.NamedTemporaryFile( + suffix="_shark_ts", prefix="temp_ts_" + ) ts_g.save(temp.name) new_ts = torch.jit.load(temp.name) self.training_graph = new_ts diff --git a/shark/examples/shark_eager/dynamo_demo.py b/shark/examples/shark_eager/dynamo_demo.py index b355c6d0..a4cf9ca9 100644 --- a/shark/examples/shark_eager/dynamo_demo.py +++ b/shark/examples/shark_eager/dynamo_demo.py @@ -8,7 +8,9 @@ try: from torchdynamo.optimizations.backends import create_backend from torchdynamo.optimizations.subgraph import SubGraph except ModuleNotFoundError: - print("Please install TorchDynamo using pip install git+https://github.com/pytorch/torchdynamo") + print( + "Please install TorchDynamo using pip install git+https://github.com/pytorch/torchdynamo" + ) exit() NUM_ITERS = 10 @@ -24,7 +26,9 @@ def __torch_mlir(fx_graph, *args, **kwargs): for node in fx_g.graph.nodes: if node.op == "output": - assert len(node.args) == 1, "Output node must have a single argument" + assert ( + len(node.args) == 1 + ), "Output node must have a single argument" node_arg = node.args[0] if isinstance(node_arg, tuple) and len(node_arg) == 1: node.args = (node_arg[0],) @@ -41,8 +45,12 @@ def __torch_mlir(fx_graph, *args, **kwargs): if len(args) == 1 and isinstance(args[0], list): args = args[0] - linalg_module = compile(ts_graph, args, output_type=OutputType.LINALG_ON_TENSORS) - callable, _ = get_iree_compiled_module(linalg_module, "cuda", func_name="forward") + linalg_module = compile( + ts_graph, args, output_type=OutputType.LINALG_ON_TENSORS + ) + callable, _ = get_iree_compiled_module( + linalg_module, "cuda", func_name="forward" + ) def forward(*inputs): return callable(*inputs) diff --git a/shark/examples/shark_inference/CLIPModel_tf.py b/shark/examples/shark_inference/CLIPModel_tf.py index 6f3e8f8c..6690d2b1 100644 --- a/shark/examples/shark_inference/CLIPModel_tf.py +++ b/shark/examples/shark_inference/CLIPModel_tf.py @@ -9,23 +9,24 @@ from shark.shark_inference import SharkInference clip_vit_inputs = [ tf.TensorSpec(shape=[2, 7], dtype=tf.int32), tf.TensorSpec(shape=[2, 7], dtype=tf.int32), - tf.TensorSpec(shape=[1, 3, 224, 224], dtype=tf.float32) + tf.TensorSpec(shape=[1, 3, 224, 224], dtype=tf.float32), ] class CLIPModule(tf.Module): - def __init__(self): super(CLIPModule, self).__init__() self.m = TFCLIPModel.from_pretrained("openai/clip-vit-base-patch32") self.m.predict = lambda x, y, z: self.m( - input_ids=x, attention_mask=y, pixel_values=z) + input_ids=x, attention_mask=y, pixel_values=z + ) @tf.function(input_signature=clip_vit_inputs) def forward(self, input_ids, attention_mask, pixel_values): - return self.m.predict(input_ids, attention_mask, - pixel_values).logits_per_image + return self.m.predict( + input_ids, attention_mask, pixel_values + ).logits_per_image if __name__ == "__main__": @@ -35,17 +36,26 @@ if __name__ == "__main__": url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(url, stream=True).raw) - inputs = processor(text=["a photo of a cat", "a photo of a dog"], - images=image, - return_tensors="tf", - padding=True) + inputs = processor( + text=["a photo of a cat", "a photo of a dog"], + images=image, + return_tensors="tf", + padding=True, + ) shark_module = SharkInference( CLIPModule(), - (inputs["input_ids"], inputs["attention_mask"], inputs["pixel_values"])) + (inputs["input_ids"], inputs["attention_mask"], inputs["pixel_values"]), + ) shark_module.set_frontend("tensorflow") shark_module.compile() print( - shark_module.forward((inputs["input_ids"], inputs["attention_mask"], - inputs["pixel_values"]))) + shark_module.forward( + ( + inputs["input_ids"], + inputs["attention_mask"], + inputs["pixel_values"], + ) + ) + ) diff --git a/shark/examples/shark_inference/gpt2_tf.py b/shark/examples/shark_inference/gpt2_tf.py index 224448ce..d04cac2f 100644 --- a/shark/examples/shark_inference/gpt2_tf.py +++ b/shark/examples/shark_inference/gpt2_tf.py @@ -13,7 +13,6 @@ gpt2_inputs = [ class GPT2Module(tf.Module): - def __init__(self): super(GPT2Module, self).__init__() self.m = TFGPT2Model.from_pretrained("distilgpt2") @@ -30,9 +29,10 @@ if __name__ == "__main__": tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2") text = "I love the distilled version of models." - inputs = tokenizer(text, return_tensors='tf') + inputs = tokenizer(text, return_tensors="tf") shark_module = SharkInference( - GPT2Module(), (inputs["input_ids"], inputs["attention_mask"])) + GPT2Module(), (inputs["input_ids"], inputs["attention_mask"]) + ) shark_module.set_frontend("tensorflow") shark_module.compile() print(shark_module.forward((inputs["input_ids"], inputs["attention_mask"]))) diff --git a/shark/examples/shark_inference/minilm_benchmark.py b/shark/examples/shark_inference/minilm_benchmark.py index 360d1352..3263f84d 100644 --- a/shark/examples/shark_inference/minilm_benchmark.py +++ b/shark/examples/shark_inference/minilm_benchmark.py @@ -7,17 +7,13 @@ tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") class MiniLMSequenceClassification(torch.nn.Module): - def __init__(self): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( "microsoft/MiniLM-L12-H384-uncased", # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -27,9 +23,12 @@ class MiniLMSequenceClassification(torch.nn.Module): test_input = torch.randint(2, (1, 128)) -shark_module = SharkInference(MiniLMSequenceClassification(), (test_input,), - jit_trace=True, - benchmark_mode=True) +shark_module = SharkInference( + MiniLMSequenceClassification(), + (test_input,), + jit_trace=True, + benchmark_mode=True, +) shark_module.compile() shark_module.forward((test_input,)) diff --git a/shark/examples/shark_inference/minilm_benchmark_tf.py b/shark/examples/shark_inference/minilm_benchmark_tf.py index 50d3b720..ce136bfb 100644 --- a/shark/examples/shark_inference/minilm_benchmark_tf.py +++ b/shark/examples/shark_inference/minilm_benchmark_tf.py @@ -9,21 +9,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() # Create a BERT trainer with the created network. self.m = TFBertModel.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased", from_pt=True) + "microsoft/MiniLM-L12-H384-uncased", from_pt=True + ) # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=bert_input) def forward(self, input_ids, attention_mask, token_type_ids): @@ -33,22 +34,26 @@ class BertModule(tf.Module): if __name__ == "__main__": # Prepping Data tokenizer = BertTokenizer.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased") + "microsoft/MiniLM-L12-H384-uncased" + ) text = "Replace me by any text you'd like." - encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + encoded_input = tokenizer( + text, + padding="max_length", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) + tf.convert_to_tensor(encoded_input[key]), 0 + ) - test_input = (encoded_input["input_ids"], encoded_input["attention_mask"], - encoded_input["token_type_ids"]) - shark_module = SharkInference( - BertModule(), - test_input, - benchmark_mode=True) + test_input = ( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ) + shark_module = SharkInference(BertModule(), test_input, benchmark_mode=True) shark_module.set_frontend("tensorflow") shark_module.compile() shark_module.benchmark_all(test_input) diff --git a/shark/examples/shark_inference/minilm_jit.py b/shark/examples/shark_inference/minilm_jit.py index 294b3c84..ab8c0b0f 100644 --- a/shark/examples/shark_inference/minilm_jit.py +++ b/shark/examples/shark_inference/minilm_jit.py @@ -7,17 +7,13 @@ tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") class MiniLMSequenceClassification(torch.nn.Module): - def __init__(self): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( "microsoft/MiniLM-L12-H384-uncased", # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -27,9 +23,11 @@ class MiniLMSequenceClassification(torch.nn.Module): test_input = torch.randint(2, (1, 128)).to(torch.int32) -shark_module = SharkInference(MiniLMSequenceClassification(), - (test_input, test_input, test_input), - jit_trace=True) +shark_module = SharkInference( + MiniLMSequenceClassification(), + (test_input, test_input, test_input), + jit_trace=True, +) shark_module.compile() result = shark_module.forward((test_input, test_input, test_input)) diff --git a/shark/examples/shark_inference/minilm_tf.py b/shark/examples/shark_inference/minilm_tf.py index a920819c..57e2c5dc 100644 --- a/shark/examples/shark_inference/minilm_tf.py +++ b/shark/examples/shark_inference/minilm_tf.py @@ -9,21 +9,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() # Create a BERT trainer with the created network. self.m = TFBertModel.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased", from_pt=True) + "microsoft/MiniLM-L12-H384-uncased", from_pt=True + ) # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=bert_input) def forward(self, input_ids, attention_mask, token_type_ids): @@ -33,24 +34,37 @@ class BertModule(tf.Module): if __name__ == "__main__": # Prepping Data tokenizer = BertTokenizer.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased") + "microsoft/MiniLM-L12-H384-uncased" + ) text = "Replace me by any text you'd like." - encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + encoded_input = tokenizer( + text, + padding="max_length", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) + tf.convert_to_tensor(encoded_input[key]), 0 + ) shark_module = SharkInference( BertModule(), - (encoded_input["input_ids"], encoded_input["attention_mask"], - encoded_input["token_type_ids"])) + ( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ), + ) shark_module.set_frontend("tensorflow") shark_module.compile() print( shark_module.forward( - (encoded_input["input_ids"], encoded_input["attention_mask"], - encoded_input["token_type_ids"]))) + ( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ) + ) + ) diff --git a/shark/examples/shark_inference/resnet50_script.py b/shark/examples/shark_inference/resnet50_script.py index 9f8bbd2a..fc5082ac 100644 --- a/shark/examples/shark_inference/resnet50_script.py +++ b/shark/examples/shark_inference/resnet50_script.py @@ -10,19 +10,22 @@ from shark.shark_inference import SharkInference ################################## Preprocessing inputs and model ############ def load_and_preprocess_image(url: str): headers = { - "User-Agent": - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36" + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36" } - img = Image.open(requests.get(url, headers=headers, - stream=True).raw).convert("RGB") + img = Image.open( + requests.get(url, headers=headers, stream=True).raw + ).convert("RGB") # preprocessing pipeline - preprocess = transforms.Compose([ - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - transforms.Normalize(mean=[0.485, 0.456, 0.406], - std=[0.229, 0.224, 0.225]), - ]) + preprocess = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) img_preprocessed = preprocess(img) return torch.unsqueeze(img_preprocessed, 0) @@ -44,7 +47,6 @@ def top3_possibilities(res): class Resnet50Module(torch.nn.Module): - def __init__(self): super().__init__() self.resnet = models.resnet50(pretrained=True) diff --git a/shark/examples/shark_inference/t5_tf.py b/shark/examples/shark_inference/t5_tf.py index bec01ecb..bbd1048d 100644 --- a/shark/examples/shark_inference/t5_tf.py +++ b/shark/examples/shark_inference/t5_tf.py @@ -11,12 +11,12 @@ t5_inputs = [ tf.TensorSpec(shape=[1, 10], dtype=tf.int32), ] -class T5Module(tf.Module): +class T5Module(tf.Module): def __init__(self): super(T5Module, self).__init__() self.m = TFT5Model.from_pretrained("t5-small") - self.m.predict = lambda x,y: self.m(input_ids=x, decoder_input_ids=y) + self.m.predict = lambda x, y: self.m(input_ids=x, decoder_input_ids=y) @tf.function(input_signature=t5_inputs) def forward(self, input_ids, decoder_input_ids): @@ -27,12 +27,9 @@ if __name__ == "__main__": # Prepping Data tokenizer = T5Tokenizer.from_pretrained("t5-small") text = "I love the distilled version of models." - inputs = tokenizer( - text, return_tensors="tf" - ).input_ids + inputs = tokenizer(text, return_tensors="tf").input_ids - shark_module = SharkInference( - T5Module(), (inputs, inputs)) + shark_module = SharkInference(T5Module(), (inputs, inputs)) shark_module.set_frontend("tensorflow") shark_module.compile() - print(shark_module.forward((inputs,inputs))) + print(shark_module.forward((inputs, inputs))) diff --git a/shark/examples/shark_inference/torch_vision_models_script.py b/shark/examples/shark_inference/torch_vision_models_script.py index 55ec9f1f..ae66883b 100644 --- a/shark/examples/shark_inference/torch_vision_models_script.py +++ b/shark/examples/shark_inference/torch_vision_models_script.py @@ -4,7 +4,6 @@ from shark.shark_inference import SharkInference class VisionModule(torch.nn.Module): - def __init__(self, model): super().__init__() self.model = model diff --git a/shark/examples/shark_inference/unet_script.py b/shark/examples/shark_inference/unet_script.py index 02ac065a..36804ed6 100644 --- a/shark/examples/shark_inference/unet_script.py +++ b/shark/examples/shark_inference/unet_script.py @@ -4,7 +4,6 @@ from shark_runner import SharkInference # Currently not supported aten.transpose_conv2d missing. class UnetModule(torch.nn.Module): - def __init__(self): super().__init__() self.model = torch.hub.load( diff --git a/shark/examples/shark_training/bert_training.py b/shark/examples/shark_training/bert_training.py index 8e2918d0..f5587978 100644 --- a/shark/examples/shark_training/bert_training.py +++ b/shark/examples/shark_training/bert_training.py @@ -5,17 +5,13 @@ from shark.shark_runner import SharkTrainer class MiniLMSequenceClassification(torch.nn.Module): - def __init__(self): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( "microsoft/MiniLM-L12-H384-uncased", # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -37,8 +33,9 @@ inp = (torch.randint(2, (1, 128)),) def forward(params, buffers, args): params_and_buffers = {**params, **buffers} - _stateless.functional_call(mod, params_and_buffers, args, - {}).sum().backward() + _stateless.functional_call( + mod, params_and_buffers, args, {} + ).sum().backward() optim = torch.optim.SGD(get_sorted_params(params), lr=0.01) # optim.load_state_dict(optim_state) optim.step() diff --git a/shark/examples/shark_training/bert_training_load_tf.py b/shark/examples/shark_training/bert_training_load_tf.py index 3d93915c..ae214c4d 100644 --- a/shark/examples/shark_training/bert_training_load_tf.py +++ b/shark/examples/shark_training/bert_training_load_tf.py @@ -11,7 +11,8 @@ parser.add_argument( "--download_mlir_path", type=str, default="bert_tf_training.mlir", - help="Specifies path to target mlir file that will be loaded.") + help="Specifies path to target mlir file that will be loaded.", +) load_args, unknown = parser.parse_known_args() tf.random.set_seed(0) @@ -25,20 +26,30 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), ] file_link = "https://storage.googleapis.com/shark_tank/users/stanley/bert_tf_training.mlir" response = request.urlretrieve(file_link, load_args.download_mlir_path) - sample_input_tensors = [tf.convert_to_tensor(val, dtype=tf.int32) for val in predict_sample_input] + sample_input_tensors = [ + tf.convert_to_tensor(val, dtype=tf.int32) + for val in predict_sample_input + ] num_iter = 10 if not os.path.isfile(load_args.download_mlir_path): - raise ValueError(f"Tried looking for target mlir in {load_args.download_mlir_path}, but cannot be found.") + raise ValueError( + f"Tried looking for target mlir in {load_args.download_mlir_path}, but cannot be found." + ) with open(load_args.download_mlir_path, "rb") as input_file: bert_mlir = input_file.read() shark_module = SharkTrainer( bert_mlir, - (sample_input_tensors, - tf.convert_to_tensor(np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32))) + ( + sample_input_tensors, + tf.convert_to_tensor( + np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32 + ), + ), + ) shark_module.set_frontend("mhlo") shark_module.compile() start = time.time() diff --git a/shark/examples/shark_training/bert_training_tf.py b/shark/examples/shark_training/bert_training_tf.py index a600f419..18a6f43d 100644 --- a/shark/examples/shark_training/bert_training_tf.py +++ b/shark/examples/shark_training/bert_training_tf.py @@ -25,17 +25,17 @@ bert_input = [ class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=2, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -46,10 +46,12 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - bert_input, # inputs - tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels - ]) + @tf.function( + input_signature=[ + bert_input, # inputs + tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels + ] + ) def forward(self, inputs, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -67,14 +69,22 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), + ] + sample_input_tensors = [ + tf.convert_to_tensor(val, dtype=tf.int32) + for val in predict_sample_input ] - sample_input_tensors = [tf.convert_to_tensor(val, dtype=tf.int32) for val in predict_sample_input] num_iter = 10 shark_module = SharkTrainer( BertModule(), - (sample_input_tensors, - tf.convert_to_tensor(np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32))) + ( + sample_input_tensors, + tf.convert_to_tensor( + np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32 + ), + ), + ) shark_module.set_frontend("tensorflow") shark_module.compile() start = time.time() diff --git a/shark/examples/shark_training/neural_net_training.py b/shark/examples/shark_training/neural_net_training.py index 1fc8ca37..a28397e0 100644 --- a/shark/examples/shark_training/neural_net_training.py +++ b/shark/examples/shark_training/neural_net_training.py @@ -4,7 +4,6 @@ from shark.shark_trainer import SharkTrainer class Foo(torch.nn.Module): - def __init__(self): super(Foo, self).__init__() self.l1 = torch.nn.Linear(10, 16) @@ -28,8 +27,9 @@ def get_sorted_params(named_params): def forward(params, buffers, args): params_and_buffers = {**params, **buffers} - _stateless.functional_call(mod, params_and_buffers, args, - {}).sum().backward() + _stateless.functional_call( + mod, params_and_buffers, args, {} + ).sum().backward() optim = torch.optim.SGD(get_sorted_params(params), lr=0.01) optim.step() return params, buffers diff --git a/shark/iree_eager_backend.py b/shark/iree_eager_backend.py index 4e272fee..068d06ac 100644 --- a/shark/iree_eager_backend.py +++ b/shark/iree_eager_backend.py @@ -28,9 +28,14 @@ from torch_mlir.eager_mode.torch_mlir_eager_backend import ( TorchMLIREagerBackend, TensorMetaData, ) -from torch_mlir_e2e_test.eager_backends.refbackend import NUMPY_TO_TORCH_DTYPE_DICT +from torch_mlir_e2e_test.eager_backends.refbackend import ( + NUMPY_TO_TORCH_DTYPE_DICT, +) -from shark.iree_utils import get_iree_compiled_module, IREE_DEVICE_MAP +from shark.iree_utils.compile_utils import ( + get_iree_compiled_module, + IREE_DEVICE_MAP, +) class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend): @@ -46,15 +51,16 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend): self.iree_device_str = IREE_DEVICE_MAP[device] self.config = ireert.Config(self.iree_device_str) - def get_torch_metadata(self, tensor: DeviceArray, - kwargs: Dict[str, Any]) -> TensorMetaData: + def get_torch_metadata( + self, tensor: DeviceArray, kwargs: Dict[str, Any] + ) -> TensorMetaData: return TensorMetaData( size=tensor.shape, dtype=NUMPY_TO_TORCH_DTYPE_DICT[tensor.dtype.type], device=torch.device(self.torch_device_str), requires_grad=tensor.dtype.type - in {np.float, np.float32, np.float64} and - kwargs.get("requires_grad", False), + in {np.float, np.float32, np.float64} + and kwargs.get("requires_grad", False), ) def compile(self, imported_module: Module): @@ -64,9 +70,9 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend): "torch-function-to-torch-backend-pipeline,torch-backend-to-linalg-on-tensors-backend-pipeline", "EagerMode", ) - callable, _ = get_iree_compiled_module(imported_module, - self.iree_device_str, - func_name=fn_name) + callable, _ = get_iree_compiled_module( + imported_module, self.iree_device_str, func_name=fn_name + ) return callable def copy_into(self, dst, src): @@ -76,6 +82,7 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend): def transfer_from_device_to_torch(self, e): return torch.from_numpy(e.to_host()) - def transfer_from_torch_to_device(self, - tensor: torch.Tensor) -> DeviceArray: + def transfer_from_torch_to_device( + self, tensor: torch.Tensor + ) -> DeviceArray: return iree.runtime.asdevicearray(self.config.device, tensor.numpy()) diff --git a/shark/iree_utils.py b/shark/iree_utils.py deleted file mode 100644 index d1f86f16..00000000 --- a/shark/iree_utils.py +++ /dev/null @@ -1,381 +0,0 @@ -# Copyright 2020 The Nod Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import iree.runtime as ireert -import iree.runtime.scripts.iree_benchmark_module as benchmark_module -import iree.compiler as ireec -from shark.torch_mlir_utils import get_module_name_for_asm_dump -from shark.cuda_utils import get_cuda_sm_cc -from shark.model_annotation import * -import subprocess -import numpy as np -import os -import re -import sys - -IREE_DEVICE_MAP = { - "cpu": "local-task", - "gpu": "cuda", - "cuda": "cuda", - "vulkan": "vulkan", - "metal": "vulkan", - "rocm": "rocm" -} - -IREE_TARGET_MAP = { - "cpu": "dylib", - "gpu": "cuda", - "cuda": "cuda", - "vulkan": "vulkan", - "metal": "vulkan", - "rocm": "rocm" -} - -UNIT_TO_SECOND_MAP = {"ms": 0.001, "s": 1} - - -def check_device_drivers(device): - """Checks necessary drivers present for gpu and vulkan devices""" - if (device in ["gpu", "cuda"]): - try: - subprocess.check_output('nvidia-smi') - except Exception: - return True - elif (device in ["metal", "vulkan"]): - try: - subprocess.check_output('vulkaninfo') - except Exception: - return True - elif (device == "cpu"): - return False - # Unknown device. - else: - return True - - return False - - -def get_iree_cpu_args(): - find_triple_cmd = "uname -s -m" - os_name, proc_name = subprocess.run( - find_triple_cmd, shell=True, stdout=subprocess.PIPE, - check=True).stdout.decode('utf-8').split() - if os_name == "Darwin": - find_kernel_version_cmd = "uname -r" - kernel_version = subprocess.run(find_kernel_version_cmd, - shell=True, - stdout=subprocess.PIPE, - check=True).stdout.decode('utf-8') - target_triple = f"{proc_name}-apple-darwin{kernel_version}" - elif os_name == "Linux": - target_triple = f"{proc_name}-linux-gnu" - else: - error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)" - raise Exception(error_message) - print(f"Target triple found:{target_triple}") - return [f"-iree-llvm-target-triple={target_triple}"] - - -def get_iree_gpu_args(): - ireert.flags.FUNCTION_INPUT_VALIDATION = False - ireert.flags.parse_flags("--cuda_allow_inline_execution") - sm_arch = get_cuda_sm_cc() - if sm_arch in ['sm_70', 'sm_72', 'sm_75', 'sm_80', 'sm_84', 'sm_86']: - return [ - "--iree-hal-cuda-disable-loop-nounroll-wa", - f"--iree-hal-cuda-llvm-target-arch={sm_arch}" - ] - else: - return ["--iree-hal-cuda-disable-loop-nounroll-wa"] - - -def get_vulkan_triple_flag(): - vulkan_device_cmd = "vulkaninfo | grep deviceName | awk \'END{{print $NF}}\'" - vulkan_device = run_cmd(vulkan_device_cmd).strip() - if vulkan_device == "M1": - print("Found Apple Device. Using m1-moltenvk-macos") - return "-iree-vulkan-target-triple=m1-moltenvk-macos" - elif vulkan_device == "A100-SXM4-40GB": - print("Found Nvidia Device. Using ampere-rtx3080-linux") - return "-iree-vulkan-target-triple=ampere-rtx3080-linux" - else: - print( - "Optimized kernel for your target device is not added yet. Contact SHARK Admin on discord[https://discord.com/invite/RUqY2h2s9u] or pull up an issue." - ) - return None - - -def get_iree_vulkan_args(): - #vulkan_flag = ["--iree-flow-demote-i64-to-i32"] - vulkan_flag = [] - vulkan_triple_flag = get_vulkan_triple_flag() - if vulkan_triple_flag is not None: - vulkan_flag.append(vulkan_triple_flag) - return vulkan_flag - - -def get_iree_device_args(device): - if device == "cpu": - return get_iree_cpu_args() - if device in ["gpu", "cuda"]: - return get_iree_gpu_args() - if device in ["metal", "vulkan"]: - return get_iree_vulkan_args() - return [] - - -def get_iree_frontend_args(frontend): - if frontend in ["torch", "pytorch", "linalg"]: - return ["--iree-llvm-target-cpu-features=host"] - elif frontend in ["tensorflow", "tf", "mhlo"]: - return [ - "--iree-llvm-target-cpu-features=host", - "--iree-mhlo-demote-i64-to-i32=false", - "--iree-flow-demote-i64-to-i32" - ] - else: - # Frontend not found. - return [] - -def get_iree_common_args(): - return [ - "--iree-stream-resource-index-bits=64", - "--iree-vm-target-index-bits=64" - ] - -def compile_module_to_flatbuffer(module, device, frontend, func_name, - model_config_path): - # Setup Compile arguments wrt to frontends. - input_type = "" - args = get_iree_frontend_args(frontend) - args += get_iree_device_args(device) - args += get_iree_common_args() - - if frontend in ["tensorflow", "tf"]: - input_type = "mhlo" - elif frontend in ["mhlo", "tosa"]: - input_type = frontend - elif frontend in ["tflite", "tflite-tosa"]: - input_type = "tosa" - - # Annotate the input module with the configs - if model_config_path != None: - # Currently tuned model only works on tf frontend - if frontend in ["tensorflow", "tf"]: - input_module = module.decode('utf-8') - elif frontend in ["pytorch", "torch"]: - input_module = module.operation.get_asm() - with create_context() as ctx: - module = model_annotation(ctx, - input_contents=input_module, - config_path=model_config_path) - module = str(module) - - # Compile according to the input type, else just try compiling. - if input_type not in ["mhlo", "tosa"]: - module = str(module) - if input_type != "": - # Currently for MHLO/TOSA. - flatbuffer_blob = ireec.compile_str( - module, - target_backends=[IREE_TARGET_MAP[device]], - extra_args=args, - input_type=input_type) - else: - # Currently for Torch. - flatbuffer_blob = ireec.compile_str( - str(module), - target_backends=[IREE_TARGET_MAP[device]], - extra_args=args) - return flatbuffer_blob - - -def get_iree_module(flatbuffer_blob, device, func_name): - vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob) - config = ireert.Config(IREE_DEVICE_MAP[device]) - ctx = ireert.SystemContext(config=config) - ctx.add_vm_module(vm_module) - ModuleCompiled = ctx.modules.module[func_name] - return ModuleCompiled, config - - -def get_iree_compiled_module(module, - device: str, - frontend: str = "torch", - func_name: str = "forward", - model_config_path: str = None): - """Given a module returns the compiled .vmfb and configs""" - flatbuffer_blob = compile_module_to_flatbuffer(module, device, frontend, - func_name, model_config_path) - return get_iree_module(flatbuffer_blob, device, func_name) - - -def export_iree_module_to_vmfb(module, - device: str, - directory: str, - frontend: str = "torch", - func_name: str = "forward", - model_config_path: str = None): - flatbuffer_blob = compile_module_to_flatbuffer(module, device, frontend, - func_name, model_config_path) - module_name = f"{frontend}_{func_name}_{device}" - filename = os.path.join(directory, module_name + ".vmfb") - print(f"Saved vmfb in {filename}.") - with open(filename, 'wb') as f: - f.write(flatbuffer_blob) - return filename - - -def export_module_to_mlir_file(module, frontend, directory: str): - mlir_str = module - if frontend in ["tensorflow", "tf", "mhlo", "tflite"]: - mlir_str = module.decode('utf-8') - elif frontend in ["pytorch", "torch"]: - mlir_str = module.operation.get_asm() - filename = os.path.join(directory, "model.mlir") - with open(filename, 'w') as f: - f.write(mlir_str) - print(f"Saved mlir in {filename}.") - return filename - - -def get_results(compiled_vm, input, config, frontend="torch"): - """Runs a .vmfb file given inputs and config and returns output.""" - device_inputs = input - if frontend in ["torch", "pytorch"]: - device_inputs = [ireert.asdevicearray(config.device, a) for a in input] - if frontend in ["tensorflow", "tf", "tflite", "tflite-tosa"]: - device_inputs = [] - for a in input: - if (isinstance(a, list)): - device_inputs.append([ - ireert.asdevicearray(config.device, val, dtype=val.dtype) - for val in a - ]) - else: - device_inputs.append(ireert.asdevicearray(config.device, a)) - result = compiled_vm(*device_inputs) - result_tensors = [] - if (isinstance(result, tuple)): - for val in result: - result_tensors.append(np.copy(np.asarray(val, val.dtype))) - return result_tensors - elif (isinstance(result, dict)): - data = list(result.items()) - res = np.array(data, dtype=object) - return np.copy(res) - else: - return np.copy(np.asarray(result, dtype=result.dtype)) - - -######### Benchmark Related Tools ########### - - -def tensor_to_type_str(input_tensors: tuple, frontend: str): - """ - Input: A tuple of input tensors i.e tuple(torch.tensor) - Output: list of string that represent mlir types (i.e 1x24xf64) - # TODO: Support more than floats, and ints - """ - list_of_type = [] - for input_tensor in input_tensors: - if (isinstance(input_tensor, tuple)): - for val in input_tensor: - type_string = "x".join([str(dim) for dim in val.shape]) - if frontend in ["torch", "pytorch"]: - dtype_string = str(val.dtype).replace("torch.", "") - elif frontend in ["tensorflow", "tf"]: - dtype = val.dtype - dtype_string = re.findall('\'[^"]*\'', - str(dtype))[0].replace("\'", "") - regex_split = re.compile("([a-zA-Z]+)([0-9]+)") - match = regex_split.match(dtype_string) - mlir_type_string = str(match.group(1)[0]) + str(match.group(2)) - type_string += f"x{mlir_type_string}" - list_of_type.append(type_string) - - else: - type_string = "x".join([str(dim) for dim in input_tensor.shape]) - if frontend in ["torch", "pytorch"]: - dtype_string = str(input_tensor.dtype).replace("torch.", "") - elif frontend in ["tensorflow", "tf"]: - dtype = input_tensor.dtype - dtype_string = re.findall('\'[^"]*\'', - str(dtype))[0].replace("\'", "") - regex_split = re.compile("([a-zA-Z]+)([0-9]+)") - match = regex_split.match(dtype_string) - mlir_type_string = str(match.group(1)[0]) + str(match.group(2)) - type_string += f"x{mlir_type_string}" - list_of_type.append(type_string) - return list_of_type - - -def build_benchmark_args(input_file: str, - device: str, - input_tensors: tuple, - frontend: str, - training=False): - """ - Inputs: input_file leading to vmfb, input_tensor to function, target device, and whether it is training or not. - Outputs: string that execute benchmark-module on target model. - """ - path = benchmark_module.__path__[0] - benchmarker_path = os.path.join(path, "..", "..", "iree-benchmark-module") - benchmark_cl = [benchmarker_path, f"--module_file={input_file}"] - fn_name = "forward" - if training == True: - # TODO: Replace name of train with actual train fn name. - fn_name = "train" - benchmark_cl.append(f"--entry_function={fn_name}") - benchmark_cl.append(f"--device={IREE_DEVICE_MAP[device]}") - mlir_input_types = tensor_to_type_str(input_tensors, frontend) - for mlir_input in mlir_input_types: - benchmark_cl.append(f"--function_input={mlir_input}") - time_extractor = "| awk \'END{{print $2 $3}}\'" - benchmark_cl.append(time_extractor) - return benchmark_cl - - -def run_cmd(cmd): - """ - Inputs: cli command string. - """ - try: - result = subprocess.run(cmd, - shell=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - check=True) - result_str = result.stdout.decode() - return result_str - except Exception: - sys.exit("Exiting program due to error running:", cmd) - - -def run_benchmark_module(benchmark_cl): - """ - Run benchmark command, extract result and return iteration/seconds. - - Input: benchmark command. - """ - benchmark_path = benchmark_cl[0] - assert os.path.exists( - benchmark_path - ), "Cannot find benchmark_module, Please contact SHARK maintainer on discord." - bench_result = run_cmd(' '.join(benchmark_cl)) - regex_split = re.compile("([0-9]+[.]*[0-9]*)([a-zA-Z]+)") - match = regex_split.match(bench_result) - time = float(match.group(1)) - unit = match.group(2) - return 1.0 / (time * UNIT_TO_SECOND_MAP[unit]) diff --git a/shark/iree_utils/__init__.py b/shark/iree_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/shark/iree_utils/_common.py b/shark/iree_utils/_common.py new file mode 100644 index 00000000..42ffa0fc --- /dev/null +++ b/shark/iree_utils/_common.py @@ -0,0 +1,77 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +## Common utilities to be shared by iree utilities. + +import os +import sys +import subprocess + + +def run_cmd(cmd): + """ + Inputs: cli command string. + """ + try: + result = subprocess.run( + cmd, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + result_str = result.stdout.decode() + return result_str + except Exception: + sys.exit("Exiting program due to error running:", cmd) + + +IREE_DEVICE_MAP = { + "cpu": "local-task", + "gpu": "cuda", + "cuda": "cuda", + "vulkan": "vulkan", + "metal": "vulkan", + "rocm": "rocm", +} + +IREE_TARGET_MAP = { + "cpu": "dylib", + "gpu": "cuda", + "cuda": "cuda", + "vulkan": "vulkan", + "metal": "vulkan", + "rocm": "rocm", +} + +# Finds whether the required drivers are installed for the given device. +def check_device_drivers(device): + """Checks necessary drivers present for gpu and vulkan devices""" + if device in ["gpu", "cuda"]: + try: + subprocess.check_output("nvidia-smi") + except Exception: + return True + elif device in ["metal", "vulkan"]: + try: + subprocess.check_output("vulkaninfo") + except Exception: + return True + elif device == "cpu": + return False + # Unknown device. + else: + return True + + return False diff --git a/shark/iree_utils/benchmark_utils.py b/shark/iree_utils/benchmark_utils.py new file mode 100644 index 00000000..10e22bb4 --- /dev/null +++ b/shark/iree_utils/benchmark_utils.py @@ -0,0 +1,94 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import iree.runtime.scripts.iree_benchmark_module as benchmark_module +from shark.iree_utils._common import run_cmd, IREE_DEVICE_MAP +import numpy as np +import os +import re + +UNIT_TO_SECOND_MAP = {"ms": 0.001, "s": 1} + + +def tensor_to_type_str(input_tensors: tuple, frontend: str): + """ + Input: A tuple of input tensors i.e tuple(torch.tensor) + Output: list of string that represent mlir types (i.e 1x24xf64) + # TODO: Support more than floats, and ints + """ + list_of_type = [] + for input_tensor in input_tensors: + type_string = "x".join([str(dim) for dim in input_tensor.shape]) + if frontend in ["torch", "pytorch"]: + dtype_string = str(input_tensor.dtype).replace("torch.", "") + elif frontend in ["tensorflow", "tf"]: + dtype = input_tensor.dtype + dtype_string = re.findall("'[^\"]*'", str(dtype))[0].replace( + "'", "" + ) + regex_split = re.compile("([a-zA-Z]+)([0-9]+)") + match = regex_split.match(dtype_string) + mlir_type_string = str(match.group(1)[0]) + str(match.group(2)) + type_string += f"x{mlir_type_string}" + list_of_type.append(type_string) + return list_of_type + + +def build_benchmark_args( + input_file: str, + device: str, + input_tensors: tuple, + frontend: str, + training=False, +): + """ + Inputs: input_file leading to vmfb, input_tensor to function, target device, + and whether it is training or not. + Outputs: string that execute benchmark-module on target model. + """ + path = benchmark_module.__path__[0] + benchmarker_path = os.path.join(path, "..", "..", "iree-benchmark-module") + benchmark_cl = [benchmarker_path, f"--module_file={input_file}"] + # TODO: The function named can be passed as one of the args. + fn_name = "forward" + if training == True: + # TODO: Replace name of train with actual train fn name. + fn_name = "train" + benchmark_cl.append(f"--entry_function={fn_name}") + benchmark_cl.append(f"--device={IREE_DEVICE_MAP[device]}") + mlir_input_types = tensor_to_type_str(input_tensors, frontend) + for mlir_input in mlir_input_types: + benchmark_cl.append(f"--function_input={mlir_input}") + time_extractor = "| awk 'END{{print $2 $3}}'" + benchmark_cl.append(time_extractor) + return benchmark_cl + + +def run_benchmark_module(benchmark_cl): + """ + Run benchmark command, extract result and return iteration/seconds. + + # TODO: Add an example of the benchmark command. + Input: benchmark command. + """ + benchmark_path = benchmark_cl[0] + assert os.path.exists( + benchmark_path + ), "Cannot find benchmark_module, Please contact SHARK maintainer on discord." + bench_result = run_cmd(" ".join(benchmark_cl)) + regex_split = re.compile("([0-9]+[.]*[0-9]*)([a-zA-Z]+)") + match = regex_split.match(bench_result) + time = float(match.group(1)) + unit = match.group(2) + return 1.0 / (time * UNIT_TO_SECOND_MAP[unit]) diff --git a/shark/iree_utils/compile_utils.py b/shark/iree_utils/compile_utils.py new file mode 100644 index 00000000..48c6b28b --- /dev/null +++ b/shark/iree_utils/compile_utils.py @@ -0,0 +1,188 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import iree.runtime as ireert +import iree.compiler as ireec +from shark.iree_utils._common import IREE_DEVICE_MAP, IREE_TARGET_MAP +from shark.model_annotation import * +import numpy as np +import os + +# Get the iree-compile arguments given device. +def get_iree_device_args(device): + if device == "cpu": + from shark.iree_utils.cpu_utils import get_iree_cpu_args + + return get_iree_cpu_args() + if device in ["gpu", "cuda"]: + from shark.iree_utils.gpu_utils import get_iree_gpu_args + + return get_iree_gpu_args() + if device in ["metal", "vulkan"]: + from shark.iree_utils.vulkan_utils import get_iree_vulkan_args + + return get_iree_vulkan_args() + return [] + + +# Get the iree-compiler arguments given frontend. +def get_iree_frontend_args(frontend): + if frontend in ["torch", "pytorch", "linalg"]: + return ["--iree-llvm-target-cpu-features=host"] + elif frontend in ["tensorflow", "tf", "mhlo"]: + return [ + "--iree-llvm-target-cpu-features=host", + "--iree-mhlo-demote-i64-to-i32=false", + "--iree-flow-demote-i64-to-i32", + ] + else: + # Frontend not found. + return [] + + +# Common args to be used given any frontend or device. +def get_iree_common_args(): + return [ + "--iree-stream-resource-index-bits=64", + "--iree-vm-target-index-bits=64", + ] + + +def compile_module_to_flatbuffer( + module, device, frontend, func_name, model_config_path +): + # Setup Compile arguments wrt to frontends. + input_type = "" + args = get_iree_frontend_args(frontend) + args += get_iree_device_args(device) + args += get_iree_common_args() + + if frontend in ["tensorflow", "tf"]: + input_type = "mhlo" + elif frontend in ["mhlo", "tosa"]: + input_type = frontend + elif frontend in ["tflite", "tflite-tosa"]: + input_type = "tosa" + + # TODO: make it simpler. + # Compile according to the input type, else just try compiling. + if input_type not in ["mhlo", "tosa"]: + module = str(module) + if input_type != "": + # Currently for MHLO/TOSA. + flatbuffer_blob = ireec.compile_str( + module, + target_backends=[IREE_TARGET_MAP[device]], + extra_args=args, + input_type=input_type, + ) + else: + # Currently for Torch. + flatbuffer_blob = ireec.compile_str( + str(module), + target_backends=[IREE_TARGET_MAP[device]], + extra_args=args, + ) + + return flatbuffer_blob + + +def get_iree_module(flatbuffer_blob, device, func_name): + # Returns the compiled module and the configs. + vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob) + config = ireert.Config(IREE_DEVICE_MAP[device]) + ctx = ireert.SystemContext(config=config) + ctx.add_vm_module(vm_module) + ModuleCompiled = ctx.modules.module[func_name] + return ModuleCompiled, config + + +def get_iree_compiled_module( + module, + device: str, + frontend: str = "torch", + func_name: str = "forward", + model_config_path: str = None, +): + """Given a module returns the compiled .vmfb and configs""" + flatbuffer_blob = compile_module_to_flatbuffer( + module, device, frontend, func_name, model_config_path + ) + return get_iree_module(flatbuffer_blob, device, func_name) + + +def export_iree_module_to_vmfb( + module, + device: str, + directory: str, + frontend: str = "torch", + func_name: str = "forward", + model_config_path: str = None, +): + # Compiles the module given specs and saves it as .vmfb file. + flatbuffer_blob = compile_module_to_flatbuffer( + module, device, frontend, func_name, model_config_path + ) + module_name = f"{frontend}_{func_name}_{device}" + filename = os.path.join(directory, module_name + ".vmfb") + print(f"Saved vmfb in {filename}.") + with open(filename, "wb") as f: + f.write(flatbuffer_blob) + return filename + + +def export_module_to_mlir_file(module, frontend, directory: str): + # TODO: write proper documentation. + mlir_str = module + if frontend in ["tensorflow", "tf", "mhlo", "tflite"]: + mlir_str = module.decode("utf-8") + elif frontend in ["pytorch", "torch"]: + mlir_str = module.operation.get_asm() + filename = os.path.join(directory, "model.mlir") + with open(filename, "w") as f: + f.write(mlir_str) + print(f"Saved mlir in {filename}.") + return filename + + +def get_results(compiled_vm, input, config, frontend="torch"): + """Runs a .vmfb file given inputs and config and returns output.""" + device_inputs = input + if frontend in ["torch", "pytorch"]: + device_inputs = [ireert.asdevicearray(config.device, a) for a in input] + if frontend in ["tensorflow", "tf", "tflite", "tflite-tosa"]: + device_inputs = [] + for a in input: + if isinstance(a, list): + device_inputs.append( + [ + ireert.asdevicearray( + config.device, val, dtype=val.dtype + ) + for val in a + ] + ) + else: + device_inputs.append(ireert.asdevicearray(config.device, a)) + result = compiled_vm(*device_inputs) + result_tensors = [] + if isinstance(result, tuple): + for val in result: + result_tensors.append(np.copy(np.asarray(val, val.dtype))) + return result_tensors + elif isinstance(result, dict): + data = list(result.items()) + res = np.array(data, dtype=object) + return np.copy(res) + else: + return np.copy(np.asarray(result, dtype=result.dtype)) diff --git a/shark/iree_utils/cpu_utils.py b/shark/iree_utils/cpu_utils.py new file mode 100644 index 00000000..71b1f056 --- /dev/null +++ b/shark/iree_utils/cpu_utils.py @@ -0,0 +1,44 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# All the iree_cpu related functionalities go here. + +import subprocess + +# Get the default cpu args. +def get_iree_cpu_args(): + find_triple_cmd = "uname -s -m" + os_name, proc_name = ( + subprocess.run( + find_triple_cmd, shell=True, stdout=subprocess.PIPE, check=True + ) + .stdout.decode("utf-8") + .split() + ) + if os_name == "Darwin": + find_kernel_version_cmd = "uname -r" + kernel_version = subprocess.run( + find_kernel_version_cmd, + shell=True, + stdout=subprocess.PIPE, + check=True, + ).stdout.decode("utf-8") + target_triple = f"{proc_name}-apple-darwin{kernel_version}" + elif os_name == "Linux": + target_triple = f"{proc_name}-linux-gnu" + else: + error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)" + raise Exception(error_message) + print(f"Target triple found:{target_triple}") + return [f"-iree-llvm-target-triple={target_triple}"] diff --git a/shark/cuda_utils.py b/shark/iree_utils/gpu_utils.py similarity index 52% rename from shark/cuda_utils.py rename to shark/iree_utils/gpu_utils.py index e0202d61..2d79a731 100644 --- a/shark/cuda_utils.py +++ b/shark/iree_utils/gpu_utils.py @@ -12,9 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. +# All the iree_gpu related functionalities go here. + +import iree.runtime as ireert import ctypes -#Some constants taken from cuda.h +# Get the default gpu args given the architecture. +def get_iree_gpu_args(): + ireert.flags.FUNCTION_INPUT_VALIDATION = False + ireert.flags.parse_flags("--cuda_allow_inline_execution") + # TODO: Give the user_interface to pass the sm_arch. + sm_arch = get_cuda_sm_cc() + if sm_arch in ["sm_70", "sm_72", "sm_75", "sm_80", "sm_84", "sm_86"]: + return [ + "--iree-hal-cuda-disable-loop-nounroll-wa", + f"--iree-hal-cuda-llvm-target-arch={sm_arch}", + ] + else: + return ["--iree-hal-cuda-disable-loop-nounroll-wa"] + + +# Some constants taken from cuda.h CUDA_SUCCESS = 0 CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16 CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR = 39 @@ -23,7 +41,7 @@ CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE = 36 def get_cuda_sm_cc(): - libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll') + libnames = ("libcuda.so", "libcuda.dylib", "cuda.dll") for libname in libnames: try: cuda = ctypes.CDLL(libname) @@ -32,10 +50,10 @@ def get_cuda_sm_cc(): else: break else: - raise OSError("could not load any of: " + ' '.join(libnames)) + raise OSError("could not load any of: " + " ".join(libnames)) nGpus = ctypes.c_int() - name = b' ' * 100 + name = b" " * 100 cc_major = ctypes.c_int() cc_minor = ctypes.c_int() @@ -47,31 +65,43 @@ def get_cuda_sm_cc(): result = cuda.cuInit(0) if result != CUDA_SUCCESS: cuda.cuGetErrorString(result, ctypes.byref(error_str)) - print("cuInit failed with error code %d: %s" % - (result, error_str.value.decode())) + print( + "cuInit failed with error code %d: %s" + % (result, error_str.value.decode()) + ) return 1 result = cuda.cuDeviceGetCount(ctypes.byref(nGpus)) if result != CUDA_SUCCESS: cuda.cuGetErrorString(result, ctypes.byref(error_str)) - print("cuDeviceGetCount failed with error code %d: %s" % - (result, error_str.value.decode())) + print( + "cuDeviceGetCount failed with error code %d: %s" + % (result, error_str.value.decode()) + ) return 1 print("Found %d device(s)." % nGpus.value) for i in range(nGpus.value): result = cuda.cuDeviceGet(ctypes.byref(device), i) if result != CUDA_SUCCESS: cuda.cuGetErrorString(result, ctypes.byref(error_str)) - print("cuDeviceGet failed with error code %d: %s" % - (result, error_str.value.decode())) + print( + "cuDeviceGet failed with error code %d: %s" + % (result, error_str.value.decode()) + ) return 1 print("Device: %d" % i) - if cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), - device) == CUDA_SUCCESS: - print(" Name: %s" % (name.split(b'\0', 1)[0].decode())) - if cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), - ctypes.byref(cc_minor), - device) == CUDA_SUCCESS: - print(" Compute Capability: %d.%d" % - (cc_major.value, cc_minor.value)) + if ( + cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device) + == CUDA_SUCCESS + ): + print(" Name: %s" % (name.split(b"\0", 1)[0].decode())) + if ( + cuda.cuDeviceComputeCapability( + ctypes.byref(cc_major), ctypes.byref(cc_minor), device + ) + == CUDA_SUCCESS + ): + print( + " Compute Capability: %d.%d" % (cc_major.value, cc_minor.value) + ) sm = f"sm_{cc_major.value}{cc_minor.value}" return sm diff --git a/shark/iree_utils/vulkan_utils.py b/shark/iree_utils/vulkan_utils.py new file mode 100644 index 00000000..7ccd7448 --- /dev/null +++ b/shark/iree_utils/vulkan_utils.py @@ -0,0 +1,44 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# All the iree_vulkan related functionalities go here. + +from shark.iree_utils._common import run_cmd + + +def get_vulkan_triple_flag(): + vulkan_device_cmd = "vulkaninfo | grep deviceName | awk 'END{{print $NF}}'" + vulkan_device = run_cmd(vulkan_device_cmd).strip() + if vulkan_device == "M1": + print("Found Apple Device. Using m1-moltenvk-macos") + return "-iree-vulkan-target-triple=m1-moltenvk-macos" + elif vulkan_device == "A100-SXM4-40GB": + print("Found Nvidia Device. Using ampere-rtx3080-linux") + return "-iree-vulkan-target-triple=ampere-rtx3080-linux" + else: + print( + """Optimized kernel for your target device is not added yet. + Contact SHARK Admin on discord[https://discord.com/invite/RUqY2h2s9u] + or pull up an issue.""" + ) + return None + + +def get_iree_vulkan_args(): + # vulkan_flag = ["--iree-flow-demote-i64-to-i32"] + vulkan_flag = [] + vulkan_triple_flag = get_vulkan_triple_flag() + if vulkan_triple_flag is not None: + vulkan_flag.append(vulkan_triple_flag) + return vulkan_flag diff --git a/shark/model_annotation.py b/shark/model_annotation.py index 2c63eaef..ab2169a8 100644 --- a/shark/model_annotation.py +++ b/shark/model_annotation.py @@ -21,7 +21,8 @@ from iree.compiler import ir from iree.compiler.transforms import ireec as ireec_trans MATMUL_OP_NAMES = set( - ["linalg.matmul", "linalg.batch_matmul", "mhlo.dot", "mhlo.dot_general"]) + ["linalg.matmul", "linalg.batch_matmul", "mhlo.dot", "mhlo.dot_general"] +) idx = 0 @@ -47,7 +48,8 @@ def model_annotation(ctx: ir.Context, *, input_contents: str, config_path: str): # - Disables verification (already done above) # - Writes as binary, avoiding costly unicode conversions sys.stdout.buffer.write( - module.operation.get_asm(assume_verified=True, binary=True)) + module.operation.get_asm(assume_verified=True, binary=True) + ) return module @@ -61,14 +63,21 @@ def walk_children(op: ir.Operation, configs: List[Dict]): child_op = child_op.operation if child_op.name in MATMUL_OP_NAMES: global idx - tile_sizes, pipeline, workgroup_size, \ - split_k, pipeline_depth = parse_config(configs[idx]) + ( + tile_sizes, + pipeline, + workgroup_size, + split_k, + pipeline_depth, + ) = parse_config(configs[idx]) - add_compilation_info(child_op, - tile_sizes=tile_sizes, - pipeline=pipeline, - workgroup_size=workgroup_size, - pipeline_depth=pipeline_depth) + add_compilation_info( + child_op, + tile_sizes=tile_sizes, + pipeline=pipeline, + workgroup_size=workgroup_size, + pipeline_depth=pipeline_depth, + ) if split_k: add_split_k(child_op, split_k) @@ -80,8 +89,11 @@ def walk_children(op: ir.Operation, configs: List[Dict]): def parse_config(config: Dict): if config["pipeline"] == "GPU" or config["pipeline"] == "GPU_TENSORCORE": - pipeline = "LLVMGPUMatmulSimt" if config[ - "pipeline"] == "GPU" else "LLVMGPUMatmulTensorCore" + pipeline = ( + "LLVMGPUMatmulSimt" + if config["pipeline"] == "GPU" + else "LLVMGPUMatmulTensorCore" + ) tile_sizes = [config["work_group_tile_sizes"]] workgroup_size = config["work_group_sizes"] try: @@ -95,8 +107,9 @@ def parse_config(config: Dict): else: pipeline = config["pipeline"] tile_sizes = [ - config["work_group_tile_sizes"], config["l1_tile_sizes"], - config["vector_tile_sizes"] + config["work_group_tile_sizes"], + config["l1_tile_sizes"], + config["vector_tile_sizes"], ] workgroup_size = [] split_k = None @@ -104,9 +117,13 @@ def parse_config(config: Dict): return tile_sizes, pipeline, workgroup_size, split_k, pipeline_depth -def add_compilation_info(op: ir.Operation, tile_sizes: List[List[int]], - pipeline: str, workgroup_size: List[int], - pipeline_depth: int): +def add_compilation_info( + op: ir.Operation, + tile_sizes: List[List[int]], + pipeline: str, + workgroup_size: List[int], + pipeline_depth: int, +): # We don't have a Python binding for CompilationInfo, so we just parse # its string form. if pipeline_depth: @@ -114,13 +131,15 @@ def add_compilation_info(op: ir.Operation, tile_sizes: List[List[int]], f"#iree_codegen.compilation_info<" f"lowering_config = , " f"translation_info = <{pipeline} pipeline_depth = {pipeline_depth}>, " - f"workgroup_size = {repr(workgroup_size)}>") + f"workgroup_size = {repr(workgroup_size)}>" + ) else: attr = ir.Attribute.parse( f"#iree_codegen.compilation_info<" f"lowering_config = , " f"translation_info = <{pipeline}>, " - f"workgroup_size = {repr(workgroup_size)}>") + f"workgroup_size = {repr(workgroup_size)}>" + ) op.attributes["compilation_info"] = attr @@ -138,6 +157,6 @@ def create_context() -> ir.Context: if __name__ == "__main__": with create_context() as ctx: - model_annotation(ctx, - input_contents=sys.argv[1], - config_path=sys.argv[2]) + model_annotation( + ctx, input_contents=sys.argv[1], config_path=sys.argv[2] + ) diff --git a/shark/parser.py b/shark/parser.py index 4713aac4..7decd7e0 100644 --- a/shark/parser.py +++ b/shark/parser.py @@ -21,7 +21,8 @@ def dir_path(path): return path else: raise argparse.ArgumentTypeError( - f"readable_dir:{path} is not a valid path") + f"readable_dir:{path} is not a valid path" + ) def dir_file(path): @@ -29,43 +30,52 @@ def dir_file(path): return path else: raise argparse.ArgumentTypeError( - f"readable_file:{path} is not a valid file") + f"readable_file:{path} is not a valid file" + ) -parser = argparse.ArgumentParser(description='SHARK runner.') +parser = argparse.ArgumentParser(description="SHARK runner.") parser.add_argument( "--device", type=str, default="cpu", - help="Device on which shark_runner runs. options are cpu, gpu, and vulkan") + help="Device on which shark_runner runs. options are cpu, gpu, and vulkan", +) parser.add_argument( "--repro_dir", - help= - "Directory to which module files will be saved for reproduction or debugging.", + help="Directory to which module files will be saved for reproduction or debugging.", type=dir_path, - default="/tmp/") -parser.add_argument("--save_mlir", - default=False, - action="store_true", - help="Saves input MLIR module to /tmp/ directory.") -parser.add_argument("--save_vmfb", - default=False, - action="store_true", - help="Saves iree .vmfb module to /tmp/ directory.") + default="/tmp/", +) +parser.add_argument( + "--save_mlir", + default=False, + action="store_true", + help="Saves input MLIR module to /tmp/ directory.", +) +parser.add_argument( + "--save_vmfb", + default=False, + action="store_true", + help="Saves iree .vmfb module to /tmp/ directory.", +) parser.add_argument( "--model_config_path", help="Directory to where the tuned model config file is located.", - default=None) + default=None, +) parser.add_argument( "--num_warmup_iterations", type=int, default=2, - help="Run the model for the specified number of warmup iterations.") + help="Run the model for the specified number of warmup iterations.", +) parser.add_argument( "--num_iterations", type=int, default=1, - help="Run the model for the specified number of iterations.") + help="Run the model for the specified number of iterations.", +) shark_args, unknown = parser.parse_known_args() diff --git a/shark/shark_benchmark_runner.py b/shark/shark_benchmark_runner.py new file mode 100644 index 00000000..ef24d0f6 --- /dev/null +++ b/shark/shark_benchmark_runner.py @@ -0,0 +1,166 @@ +# Copyright 2020 The Nod Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from shark.shark_runner import SharkRunner +from shark.iree_utils.compile_utils import export_iree_module_to_vmfb +from shark.iree_utils.benchmark_utils import ( + build_benchmark_args, + run_benchmark_module, +) +from shark.parser import shark_args +from datetime import datetime +import time +import csv +import os + + +class SharkBenchmarkRunner(SharkRunner): + # SharkRunner derived class with Benchmarking capabilities. + def __init__( + self, + model, + input: tuple, + dynamic: bool = False, + device: str = None, + jit_trace: bool = False, + from_aot: bool = False, + frontend: str = "torch", + ): + SharkRunner.__init__( + self, model, input, dynamic, device, jit_trace, from_aot, frontend + ) + if self.vmfb_file == None: + self.vmfb_file = export_iree_module_to_vmfb( + self.model, device, shark_args.repro_dir, frontend + ) + self.benchmark_cl = build_benchmark_args( + self.vmfb_file, device, input, frontend, from_aot + ) + + def benchmark_frontend(self, inputs): + if self.frontend in ["pytorch", "torch"]: + return self.benchmark_torch(inputs) + elif self.frontend in ["tensorflow", "tf"]: + return self.benchmark_tf(inputs) + + def benchmark_torch(self, inputs): + inputs = self.input if self.from_aot else inputs + inputs = inputs[0] + for i in range(shark_args.num_warmup_iterations): + self.frontend_model.forward(inputs) + + begin = time.time() + for i in range(shark_args.num_iterations): + out = self.frontend_model.forward(inputs) + if i == shark_args.num_iterations - 1: + end = time.time() + break + print( + f"Torch benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" + ) + return [ + f"{shark_args.num_iterations/(end-begin)}", + f"{((end-begin)/shark_args.num_iterations)*1000}", + ] + + def benchmark_tf(self, inputs): + for i in range(shark_args.num_warmup_iterations): + self.frontend_model.forward(*inputs) + + begin = time.time() + for i in range(shark_args.num_iterations): + out = self.frontend_model.forward(*inputs) + if i == shark_args.num_iterations - 1: + end = time.time() + break + print( + f"TF benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" + ) + return [ + f"{shark_args.num_iterations/(end-begin)}", + f"{((end-begin)/shark_args.num_iterations)*1000}", + ] + + def benchmark_c(self): + result = run_benchmark_module(self.benchmark_cl) + print(f"Shark-{self.frontend} C-benchmark:{result} iter/second") + return [f"{result}", f"{1000/result}"] + + def benchmark_python(self, inputs): + inputs = self.input if self.from_aot else inputs + input_list = [x for x in inputs] + for i in range(shark_args.num_warmup_iterations): + self.forward(input_list, self.frontend) + + begin = time.time() + for i in range(shark_args.num_iterations): + out = self.forward(input_list, self.frontend) + if i == shark_args.num_iterations - 1: + end = time.time() + print( + f"Shark-{self.frontend} Python-benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" + ) + return [ + f"{shark_args.num_iterations/(end-begin)}", + f"{((end-begin)/shark_args.num_iterations)*1000}", + ] + + def benchmark_all(self, inputs): + self.benchmark_frontend(inputs) + self.benchmark_python(inputs) + self.benchmark_c() + + def benchmark_all_csv(self, inputs, modelname, dynamic, device_str): + field_names = [ + "platform", + "model", + "dynamic", + "device", + "iter/sec", + "ms/iter", + "datetime", + ] + platforms = ["frontend", "shark_python", "shark_iree_c"] + + if not os.path.exists("bench_results.csv"): + with open("bench_results.csv", mode="w", newline="") as f: + writer = csv.writer(f) + writer.writerow(field_names) + + with open("bench_results.csv", mode="a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=field_names) + bench_result = {} + bench_result["model"] = modelname + if dynamic == True: + bench_result["dynamic"] = "True" + else: + bench_result["dynamic"] = "False" + bench_result["device"] = device_str + for p in platforms: + if p == "frontend": + bench_result["platform"] = "frontend" + bench_result["iter/sec"] = self.benchmark_frontend(inputs)[ + 0 + ] + bench_result["ms/iter"] = self.benchmark_frontend(inputs)[1] + elif p == "shark_python": + bench_result["platform"] = "shark_python" + bench_result["iter/sec"] = self.benchmark_python(inputs)[0] + bench_result["ms/iter"] = self.benchmark_python(inputs)[1] + else: + bench_result["platform"] = "shark_iree_c" + bench_result["iter/sec"] = self.benchmark_c()[0] + bench_result["ms/iter"] = self.benchmark_c()[1] + bench_result["datetime"] = str(datetime.now()) + writer.writerow(bench_result) diff --git a/shark/shark_downloader.py b/shark/shark_downloader.py index 1ea83806..6333db5f 100644 --- a/shark/shark_downloader.py +++ b/shark/shark_downloader.py @@ -22,19 +22,20 @@ input_type_to_np_dtype = { "float64": np.float64, "bool": np.bool_, "int32": np.int32, - "int64": np.int64 + "int64": np.int64, } 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-tosa", - input_json: str = "input.json", - input_type: str = "int32"): + 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-tosa", + 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 @@ -65,37 +66,52 @@ class SharkDownloader: print("load json inputs") if self.model_type in ["tflite-tosa"]: args = [] - with open(self.input_json, 'r') as f: + with open(self.input_json, "r") as f: args = json.load(f) - self.inputs = [np.asarray(arg, dtype=self.input_type) for arg in args] + self.inputs = [ + np.asarray(arg, dtype=self.input_type) for arg in args + ] else: - print("No json input required for current model. You could call setup_inputs(you_inputs).") + print( + "No json input required for current model. You could call setup_inputs(you_inputs)." + ) return self.inputs def load_mlir_model(self): if self.model_type in ["tflite-tosa"]: - workdir = os.path.join(os.path.dirname(__file__), self.local_tank_dir) + workdir = os.path.join( + os.path.dirname(__file__), self.local_tank_dir + ) os.makedirs(workdir, exist_ok=True) print(f"TMP_MODEL_DIR = {workdir}") # use model name get dir. model_name_dir = os.path.join(workdir, str(self.model_name)) if not os.path.exists(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.") + 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(model_name_dir, exist_ok=True) print(f"TMP_MODELNAME_DIR = {model_name_dir}") - mlir_url = self.tank_url + "/tflite/" + str(self.model_name) + "/"+str(self.model_name) + '_tosa.mlir' - self.mlir_file = '/'.join( - [model_name_dir, str(self.model_name) + '_tosa.mlir']) + mlir_url = ( + self.tank_url + + "/tflite/" + + str(self.model_name) + + "/" + + str(self.model_name) + + "_tosa.mlir" + ) + self.mlir_file = "/".join( + [model_name_dir, str(self.model_name) + "_tosa.mlir"] + ) if os.path.exists(self.mlir_file): print("Model has been downloaded before.", self.mlir_file) else: print("Download mlir model", mlir_url) - urllib.request.urlretrieve(mlir_url, - self.mlir_file) + urllib.request.urlretrieve(mlir_url, self.mlir_file) print("Get tosa.mlir model return") with open(self.mlir_file) as f: diff --git a/shark/shark_importer.py b/shark/shark_importer.py index 3dfce1a9..71bb588b 100644 --- a/shark/shark_importer.py +++ b/shark/shark_importer.py @@ -15,19 +15,20 @@ from shark.iree_utils import IREE_TARGET_MAP class SharkImporter: - - def __init__(self, - model_name: str=None, - model_path: str=None, - model_type: str = "tflite", - model_source_hub: str = "tfhub", - device: str = None, - dynamic: bool = False, - jit_trace: bool = False, - benchmark_mode: bool = False, - input_details=None, - output_details=None, - tank_url: str = None): + def __init__( + self, + model_name: str = None, + model_path: str = None, + model_type: str = "tflite", + model_source_hub: str = "tfhub", + device: str = None, + dynamic: bool = False, + jit_trace: bool = False, + benchmark_mode: bool = False, + input_details=None, + output_details=None, + tank_url: str = None, + ): self.model_name = model_name self.model_path = model_path self.model_type = model_type @@ -45,7 +46,9 @@ class SharkImporter: # create tmp model file directory if self.model_path is None and self.model_name is None: - print("Error. No model_path, No model name,Please input either one.") + print( + "Error. No model_path, No model name,Please input either one." + ) return if self.model_source_hub == "tfhub": @@ -56,43 +59,62 @@ class SharkImporter: print("Error, load tflite model fail") return - if (self.input_details == None) or \ - (self.output_details == None): - print("Setting up tflite interpreter to get model input details") + if (self.input_details == None) or ( + self.output_details == None + ): + print( + "Setting up tflite interpreter to get model input details" + ) self.tflite_interpreter = tf.lite.Interpreter( - model_path=self.tflite_saving_file) + model_path=self.tflite_saving_file + ) self.tflite_interpreter.allocate_tensors() # default input initialization - self.input_details, self.output_details = self.get_model_details( - ) + ( + self.input_details, + self.output_details, + ) = self.get_model_details() inputs = self.generate_inputs( - self.input_details) # device_inputs + self.input_details + ) # device_inputs self.setup_inputs(inputs) def load_tflite_model(self): print("Setting up for TMP_DIR") - tflite_workdir = os.path.join(os.path.dirname(__file__), "./../gen_shark_tank/tflite") + tflite_workdir = os.path.join( + os.path.dirname(__file__), "./../gen_shark_tank/tflite" + ) os.makedirs(tflite_workdir, exist_ok=True) print(f"TMP_TFLITE_DIR = {tflite_workdir}") # use model name get dir. - tflite_model_name_dir = os.path.join(tflite_workdir, str(self.model_name)) + tflite_model_name_dir = os.path.join( + tflite_workdir, str(self.model_name) + ) # TODO Download model from google bucket to tflite_model_name_dir by tank_url os.makedirs(tflite_model_name_dir, exist_ok=True) print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}") - self.tflite_saving_file = '/'.join( - [tflite_model_name_dir, str(self.model_name) + '_tflite.tflite']) - self.tflite_tosa_file = '/'.join( - [tflite_model_name_dir, str(self.model_name) + '_tosa.mlir']) + self.tflite_saving_file = "/".join( + [tflite_model_name_dir, str(self.model_name) + "_tflite.tflite"] + ) + self.tflite_tosa_file = "/".join( + [tflite_model_name_dir, str(self.model_name) + "_tosa.mlir"] + ) if os.path.exists(self.tflite_saving_file): - print("Local address for tflite model file Exists: ", self.tflite_saving_file) + print( + "Local address for tflite model file Exists: ", + self.tflite_saving_file, + ) else: print("No local tflite file, Download tflite model") if self.model_path is None: # get model file from tflite_model_list.csv or download from gs://bucket print("No model_path, get from tflite_model_list.csv") - tflite_model_list_path = os.path.join(os.path.dirname(__file__), "../tank/tflite/tflite_model_list.csv") + tflite_model_list_path = os.path.join( + os.path.dirname(__file__), + "../tank/tflite/tflite_model_list.csv", + ) tflite_model_list = csv.reader(open(tflite_model_list_path)) for row in tflite_model_list: if str(row[0]) == self.model_name: @@ -100,12 +122,13 @@ class SharkImporter: if self.model_path is None: print("Error, No model path find in tflite_model_list.csv") return False - urllib.request.urlretrieve(self.model_path, - self.tflite_saving_file) + urllib.request.urlretrieve(self.model_path, self.tflite_saving_file) if os.path.exists(self.tflite_tosa_file): print("Exists", self.tflite_tosa_file) else: - print("No tflite tosa.mlir, please use python generate_sharktank.py to download tosa model") + print( + "No tflite tosa.mlir, please use python generate_sharktank.py to download tosa model" + ) return True def generate_inputs(self, input_details): @@ -139,21 +162,25 @@ class SharkImporter: tosa_model = [] with open(self.tflite_tosa_file) as f: tosa_model = f.read() - self.shark_module = SharkInference(tosa_model, - self.inputs, - device=self.device, - dynamic=self.dynamic, - jit_trace=self.jit_trace) + self.shark_module = SharkInference( + tosa_model, + self.inputs, + device=self.device, + dynamic=self.dynamic, + jit_trace=self.jit_trace, + ) self.shark_module.set_frontend("tflite-tosa") self.shark_module.compile() else: # compile and run tfhub tflite print("Inference tfhub tflite model") - self.shark_module = SharkInference(self.tflite_saving_file, - self.inputs, - device=self.device, - dynamic=self.dynamic, - jit_trace=self.jit_trace) + self.shark_module = SharkInference( + self.tflite_saving_file, + self.inputs, + device=self.device, + dynamic=self.dynamic, + jit_trace=self.jit_trace, + ) self.shark_module.set_frontend("tflite") self.shark_module.compile() elif self.model_source_hub == "huggingface": diff --git a/shark/shark_inference.py b/shark/shark_inference.py index 2db1e058..5955d2cb 100644 --- a/shark/shark_inference.py +++ b/shark/shark_inference.py @@ -9,10 +9,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend from shark.parser import shark_args -from shark.shark_runner import SharkRunner, SharkBenchmarkRunner +from shark.shark_runner import SharkRunner import sys +import os # Prints to stderr. @@ -23,13 +23,15 @@ def print_err(*a): class SharkInference: """Inference API targeting pytorch, tensorflow, linalg, mhlo and tosa frontend.""" - def __init__(self, - model, - input: tuple, - device: str = None, - dynamic: bool = False, - jit_trace: bool = False, - benchmark_mode: bool = False): + def __init__( + self, + model, + input: tuple, + device: str = None, + dynamic: bool = False, + jit_trace: bool = False, + benchmark_mode: bool = False, + ): self.model = model self.input = input self.dynamic = dynamic @@ -49,27 +51,48 @@ class SharkInference: # Sets the frontend i.e `pytorch` or `tensorflow`. def set_frontend(self, frontend: str): if frontend not in [ - "pytorch", "torch", "tensorflow", "tf", "mhlo", "linalg", - "tosa", "tflite", "tflite-tosa" + "pytorch", + "torch", + "tensorflow", + "tf", + "mhlo", + "linalg", + "tosa", + "tflite", + "tflite-tosa", ]: print_err("frontend not supported.") else: self.frontend = frontend def compile(self): - # Inference do not use AOT. + # Inference do not use AOT. # TODO: Remove the from_aot arg as it's not + # needed. from_aot = False - if (self.benchmark_mode == True): - self.shark_runner = SharkBenchmarkRunner(self.model, self.input, - self.dynamic, self.device, - self.jit_trace, from_aot, - self.frontend) + if self.benchmark_mode == True: + # Only import shark_benchmark runner when needed. + from shark.shark_benchmark_runner import SharkBenchmarkRunner + + self.shark_runner = SharkBenchmarkRunner( + self.model, + self.input, + self.dynamic, + self.device, + self.jit_trace, + from_aot, + self.frontend, + ) else: - self.shark_runner = SharkRunner(self.model, self.input, - self.dynamic, self.device, - self.jit_trace, from_aot, - self.frontend, - self.model_config_path) + self.shark_runner = SharkRunner( + self.model, + self.input, + self.dynamic, + self.device, + self.jit_trace, + from_aot, + self.frontend, + self.model_config_path, + ) # inputs are considered to be np.array. def forward(self, inputs): @@ -81,6 +104,9 @@ class SharkInference: input_list = [x.numpy() for x in inputs] return self.shark_runner.forward(input_list, self.frontend) + def import_mlir(self, model_name="model", dir=os.getcwd()): + self.shark_runner.import_mlir(model_name, dir) + # Saves the .vmfb module. def save_module(self, dir=None): if dir is None: @@ -89,9 +115,10 @@ class SharkInference: ######### Benchmark Related Functions ######### def benchmark_mode(func): - def inner(self, *args, **kwargs): - assert self.benchmark_mode, "SharkRunner needs to be in benchmark mode to run benchmark methods." + assert ( + self.benchmark_mode + ), "SharkRunner needs to be in benchmark mode to run benchmark methods." return func(self, *args, **kwargs) return inner @@ -114,4 +141,6 @@ class SharkInference: @benchmark_mode def benchmark_all_csv(self, inputs, modelname, dynamic, device_str): - self.shark_runner.benchmark_all_csv(inputs, modelname, dynamic, device_str) + self.shark_runner.benchmark_all_csv( + inputs, modelname, dynamic, device_str + ) diff --git a/shark/shark_runner.py b/shark/shark_runner.py index 71cf1c0c..53136a94 100644 --- a/shark/shark_runner.py +++ b/shark/shark_runner.py @@ -20,13 +20,14 @@ from torch_mlir_e2e_test.eager_backends.refbackend import EagerModeRefBackend from shark.iree_eager_backend import EagerModeIREELinalgOnTensorsBackend from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend -from shark.iree_utils import get_results, get_iree_compiled_module, export_iree_module_to_vmfb, export_module_to_mlir_file, build_benchmark_args, run_benchmark_module -import os from shark.parser import shark_args -from tqdm import tqdm -from datetime import datetime -import time -import csv +from shark.iree_utils.compile_utils import ( + get_iree_compiled_module, + export_iree_module_to_vmfb, + export_module_to_mlir_file, + get_results, +) +import os class SharkRunner: @@ -57,24 +58,27 @@ class SharkRunner: elif self.frontend in ["pytorch", "torch"]: # get torch-mlir dialect # self.model = torch.Module + # Lowers in linalg dialect. # TODO assert - self.model = get_torch_mlir_module(self.model, input, dynamic, - jit_trace, from_aot) + # TODO tosa dialect from torch_module. + self.model = get_torch_mlir_module( + self.model, input, dynamic, jit_trace, from_aot + ) elif self.frontend in ["tensorflow", "tf"]: # get mhlo dialect # self.model = tf.Module # TODO assert - self.model = tfc.compile_module(self.model, - exported_names=[func_name], - import_only=True) + self.model = tfc.compile_module( + self.model, exported_names=[func_name], import_only=True + ) elif self.frontend in ["tflite"]: print("Setting up for IREE compiler tflite") # get tosa dialect # self.model = model.tflite # TODO assert - self.model = ireec_tflite.compile_file(self.model, - input_type="tosa", - import_only=True) + self.model = ireec_tflite.compile_file( + self.model, input_type="tosa", import_only=True + ) func_name = "main" # TODO: We can capture the .vmfb module here and later use it for saving @@ -82,29 +86,44 @@ class SharkRunner: ( self.iree_compilation_module, self.iree_config, - ) = get_iree_compiled_module(self.model, - self.device, - self.frontend, - func_name=func_name, - model_config_path=model_config_path) + ) = get_iree_compiled_module( + self.model, + self.device, + self.frontend, + func_name=func_name, + model_config_path=model_config_path, + ) # Debugging Options: if shark_args.save_mlir: - export_module_to_mlir_file(self.model, self.frontend, - shark_args.repro_dir) + export_module_to_mlir_file( + self.model, self.frontend, shark_args.repro_dir + ) if shark_args.save_vmfb: self.vmfb_file = self.save_module(shark_args.repro_dir) # All the timings and benchmarking can be done here. def forward(self, input, frontend): - return get_results(self.iree_compilation_module, input, - self.iree_config, frontend) + return get_results( + self.iree_compilation_module, input, self.iree_config, frontend + ) + + # Saves the .mlir file, can be in tosa, linalg or mhlo dialect. + # torch-mlir can export tosa or linalg dialects. + # tensorflow models get exported to mhlo dialect. + def import_mlir(self, model_name, dir): + filename = os.path.join(dir, f"{model_name}_{self.frontend}.mlir") + with open(filename, "w") as f: + f.write(self.model) + print(f"Saved mlir in {filename}.") + return filename # TODO: Instead of passing directory and having names decided by the module # , user may want to save the module with manual names. def save_module(self, dir=os.getcwd()): - return export_iree_module_to_vmfb(self.model, self.device, dir, - self.frontend) + return export_iree_module_to_vmfb( + self.model, self.device, dir, self.frontend + ) # TODO: Load a module and directly use it, we will need to set the frontend # in this case. @@ -112,150 +131,17 @@ class SharkRunner: pass +# TODO: Document shark_eager mode. class SharkEagerMode: - def __init__(self, device="cpu"): if device == "refbackend": torch_mlir_tensor.backend = EagerModeRefBackend() else: torch_mlir_tensor.backend = EagerModeIREELinalgOnTensorsBackend( - device) + device + ) self.guard = enable_torch_dispatch_mode(TorchMLIRTensor) self.guard.__enter__() def __del__(self): self.guard.__exit__(None, None, None) - - -class SharkBenchmarkRunner(SharkRunner): - # SharkRunner derived class with Benchmarking capabilities. - def __init__( - self, - model, - input: tuple, - dynamic: bool = False, - device: str = None, - jit_trace: bool = False, - from_aot: bool = False, - frontend: str = "torch", - ): - SharkRunner.__init__(self, model, input, dynamic, device, jit_trace, - from_aot, frontend) - if (self.vmfb_file == None): - self.vmfb_file = export_iree_module_to_vmfb(self.model, device, - shark_args.repro_dir, - frontend) - self.benchmark_cl = build_benchmark_args(self.vmfb_file, device, input, - frontend, from_aot) - - def benchmark_frontend(self, inputs): - if self.frontend in ["pytorch", "torch"]: - return self.benchmark_torch(inputs) - elif self.frontend in ["tensorflow", "tf"]: - return self.benchmark_tf(inputs) - - def benchmark_torch(self, inputs): - inputs = self.input if self.from_aot else inputs - inputs = inputs[0] - for i in range(shark_args.num_warmup_iterations): - self.frontend_model.forward(inputs) - - begin = time.time() - for i in range(shark_args.num_iterations): - out = self.frontend_model.forward(inputs) - if i == shark_args.num_iterations - 1: - end = time.time() - break - print( - f"Torch benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" - ) - return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"] - - def benchmark_tf(self, inputs): - for i in range(shark_args.num_warmup_iterations): - self.frontend_model.forward(*inputs) - - begin = time.time() - for i in range(shark_args.num_iterations): - out = self.frontend_model.forward(*inputs) - if i == shark_args.num_iterations - 1: - end = time.time() - break - print( - f"TF benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" - ) - return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"] - - def benchmark_c(self): - result = run_benchmark_module(self.benchmark_cl) - print(f"Shark-{self.frontend} C-benchmark:{result} iter/second") - return [f"{result}", f"{1000/result}"] - - def benchmark_python(self, inputs): - inputs = self.input if self.from_aot else inputs - input_list = [x for x in inputs] - for i in range(shark_args.num_warmup_iterations): - self.forward(input_list, self.frontend) - - begin = time.time() - for i in range(shark_args.num_iterations): - out = self.forward(input_list, self.frontend) - if i == shark_args.num_iterations - 1: - end = time.time() - print( - f"Shark-{self.frontend} Python-benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}" - ) - return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"] - - def benchmark_all(self, inputs): - self.benchmark_frontend(inputs) - self.benchmark_python(inputs) - self.benchmark_c() - - def benchmark_all_csv(self, inputs, modelname, dynamic, device_str): - field_names = [ - 'platform', - 'model', - 'dynamic', - 'device', - 'iter/sec', - 'ms/iter', - 'datetime' - ] - platforms = [ - 'frontend', - 'shark_python', - 'shark_iree_c' - ] - - if not os.path.exists('bench_results.csv'): - with open('bench_results.csv', mode='w', newline='') as f: - writer = csv.writer(f) - writer.writerow(field_names) - - with open('bench_results.csv', mode='a', newline='') as f: - writer = csv.DictWriter(f, fieldnames=field_names) - bench_result = {} - bench_result['model'] = modelname - if dynamic == True: - bench_result['dynamic'] = "True" - else: - bench_result['dynamic'] = "False" - bench_result['device'] = device_str - for p in platforms: - if p == 'frontend': - bench_result['platform'] = "frontend" - bench_result['iter/sec'] = self.benchmark_frontend(inputs)[0] - bench_result['ms/iter'] = self.benchmark_frontend(inputs)[1] - elif p == 'shark_python': - bench_result['platform'] = "shark_python" - bench_result['iter/sec'] = self.benchmark_python(inputs)[0] - bench_result['ms/iter'] = self.benchmark_python(inputs)[1] - else: - bench_result['platform'] = "shark_iree_c" - bench_result['iter/sec'] = self.benchmark_c()[0] - bench_result['ms/iter'] = self.benchmark_c()[1] - bench_result['datetime'] = str(datetime.now()) - writer.writerow(bench_result) - - diff --git a/shark/shark_trainer.py b/shark/shark_trainer.py index 2f005d3d..f9835f8d 100644 --- a/shark/shark_trainer.py +++ b/shark/shark_trainer.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend -from shark.iree_utils import get_results, get_iree_compiled_module, export_iree_module_to_vmfb from shark.parser import shark_args from shark.shark_runner import SharkRunner from shark.backward_makefx import MakeFxModule @@ -56,7 +54,13 @@ class SharkTrainer: # Sets the frontend i.e `pytorch` or `tensorflow`. def set_frontend(self, frontend: str): if frontend not in [ - "pytorch", "torch", "tensorflow", "tf", "mhlo", "linalg", "tosa" + "pytorch", + "torch", + "tensorflow", + "tf", + "mhlo", + "linalg", + "tosa", ]: print_err("frontend not supported.") else: @@ -65,22 +69,32 @@ class SharkTrainer: # Training function is needed in the case of torch_fn. def compile(self, training_fn=None): if self.frontend in ["torch", "pytorch"]: - aot_module = MakeFxModule(self.model, - tuple(self.input), - custom_inference_fn=training_fn) + aot_module = MakeFxModule( + self.model, tuple(self.input), custom_inference_fn=training_fn + ) aot_module.generate_graph() # Returns the backward graph. training_graph = aot_module.training_graph weights = self.get_torch_params() - self.shark_runner = SharkRunner(training_graph, - weights + self.input, self.dynamic, - self.device, self.jit_trace, - self.from_aot, self.frontend) + self.shark_runner = SharkRunner( + training_graph, + weights + self.input, + self.dynamic, + self.device, + self.jit_trace, + self.from_aot, + self.frontend, + ) elif self.frontend in ["tensorflow", "tf", "mhlo"]: - self.shark_runner = SharkRunner(self.model, self.input, - self.dynamic, self.device, - self.jit_trace, self.from_aot, - self.frontend) + self.shark_runner = SharkRunner( + self.model, + self.input, + self.dynamic, + self.device, + self.jit_trace, + self.from_aot, + self.frontend, + ) else: print_err("Unknown frontend") return @@ -98,8 +112,9 @@ class SharkTrainer: params = [x.numpy() for x in params] print(f"Training started for {num_iters} iterations:") for i in tqdm(range(num_iters)): - params = self.shark_runner.forward(params + self.input, - self.frontend) + params = self.shark_runner.forward( + params + self.input, self.frontend + ) return params @@ -109,15 +124,15 @@ class SharkTrainer: def _train_tf(self, num_iters): input_list = [] for x in self.input: - if (isinstance(x, list)): + if isinstance(x, list): nested_list = [] for val in x: - if (isinstance(val, np.ndarray)): + if isinstance(val, np.ndarray): nested_list.append(val) else: nested_list.append(val.numpy()) input_list.append(nested_list) - elif (isinstance(x, np.ndarray)): + elif isinstance(x, np.ndarray): input_list.append(x) else: input_list.append(x.numpy()) diff --git a/shark/tests/test_shark_importer.py b/shark/tests/test_shark_importer.py index 101ca885..9f467cfd 100644 --- a/shark/tests/test_shark_importer.py +++ b/shark/tests/test_shark_importer.py @@ -13,38 +13,48 @@ def generate_inputs(input_details): args = [] args.append( - np.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + 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"])) + 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"])) + np.zeros( + shape=input_details[2]["shape"], dtype=input_details[2]["dtype"] + ) + ) return args # A specific case can be run by commenting different cases. Runs all the test # across cpu, gpu and vulkan according to available drivers. pytest_param = pytest.mark.parametrize( - ('dynamic', 'device'), + ("dynamic", "device"), [ - pytest.param(False, 'cpu'), + pytest.param(False, "cpu"), # TODO: Language models are failing for dynamic case.. - pytest.param(True, 'cpu', marks=pytest.mark.skip), - ]) + pytest.param(True, "cpu", marks=pytest.mark.skip), + ], +) @pytest_param def test_albert(dynamic, device): - my_shark_importer = SharkImporter(model_path=model_path, - model_type="tflite", - model_source_hub="tfhub", - device=device, - dynamic=dynamic, - jit_trace=True) + my_shark_importer = SharkImporter( + model_path=model_path, + model_type="tflite", + model_source_hub="tfhub", + device=device, + dynamic=dynamic, + jit_trace=True, + ) input_details, output_details = my_shark_importer.get_model_details() inputs = generate_inputs(input_details) # device_inputs my_shark_importer.compile(inputs) diff --git a/shark/torch_mlir_utils.py b/shark/torch_mlir_utils.py index 2240a19d..f65f415d 100644 --- a/shark/torch_mlir_utils.py +++ b/shark/torch_mlir_utils.py @@ -21,8 +21,10 @@ from torch_mlir.dialects.torch.importer.jit_ir import ( ModuleBuilder, ) from torch_mlir_e2e_test.torchscript.serialization import ( - extract_serializable_annotations, apply_serializable_annotations, - SerializableTest) + extract_serializable_annotations, + apply_serializable_annotations, + SerializableTest, +) from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend @@ -38,7 +40,8 @@ def get_module_name_for_asm_dump(module): if not "torch.debug_module_name" in module.operation.attributes: return "UnnammedModule" return StringAttr( - module.operation.attributes["torch.debug_module_name"]).value + module.operation.attributes["torch.debug_module_name"] + ).value def get_input_annotations(inputs: tuple, dynamic: bool) -> list: @@ -65,8 +68,9 @@ def run_on_refbackend(torch_module, inputs): return jit_module.forward(np_inputs[0]) -def shark_jit_trace(module, input: tuple, dynamic: bool, - tracing_required: bool): +def shark_jit_trace( + module, input: tuple, dynamic: bool, tracing_required: bool +): """TODO: Include necessary documentation.""" if not tracing_required: @@ -76,21 +80,26 @@ def shark_jit_trace(module, input: tuple, dynamic: bool, actual_script = traced_module._actual_script_module export(actual_script.forward) annotate_args_decorator = annotate_args( - get_input_annotations(input, dynamic)) + get_input_annotations(input, dynamic) + ) annotate_args_decorator(actual_script.forward) module = torch.jit.script(actual_script) # TODO: remove saved annotations.pickle - torchscript_module_bytes = module.save_to_buffer({ - "annotations.pkl": - pickle.dumps(extract_serializable_annotations(module)) - }) - serializable_test = SerializableTest(unique_name="", - program=torchscript_module_bytes, - trace=None) + torchscript_module_bytes = module.save_to_buffer( + { + "annotations.pkl": pickle.dumps( + extract_serializable_annotations(module) + ) + } + ) + serializable_test = SerializableTest( + unique_name="", program=torchscript_module_bytes, trace=None + ) _extra_files = {"annotations.pkl": ""} - module = torch.jit.load(io.BytesIO(serializable_test.program), - _extra_files=_extra_files) + module = torch.jit.load( + io.BytesIO(serializable_test.program), _extra_files=_extra_files + ) # Load the pickled annotations. annotations = pickle.loads(_extra_files["annotations.pkl"]) apply_serializable_annotations(module, annotations) diff --git a/tank/conftest.py b/tank/conftest.py index 3bcd2cf6..b8ed0424 100644 --- a/tank/conftest.py +++ b/tank/conftest.py @@ -1,18 +1,26 @@ def pytest_addoption(parser): # Attaches SHARK command-line arguments to the pytest machinery. - parser.addoption("--save_mlir", - action="store_true", - default="False", - help="Pass option to save input MLIR") - parser.addoption("--save_vmfb", - action="store_true", - default="False", - help="Pass option to save IREE output .vmfb") - parser.addoption("--benchmark", - action="store_true", - default="False", - help="Pass option to benchmark and write results.csv") - parser.addoption("--save_temps", - action="store_true", - default="False", - help="Saves IREE reproduction artifacts for filing upstream issues.") + parser.addoption( + "--save_mlir", + action="store_true", + default="False", + help="Pass option to save input MLIR", + ) + parser.addoption( + "--save_vmfb", + action="store_true", + default="False", + help="Pass option to save IREE output .vmfb", + ) + parser.addoption( + "--benchmark", + action="store_true", + default="False", + help="Pass option to benchmark and write results.csv", + ) + parser.addoption( + "--save_temps", + action="store_true", + default="False", + help="Saves IREE reproduction artifacts for filing upstream issues.", + ) diff --git a/tank/model_utils.py b/tank/model_utils.py index ecc46035..24d0bc3e 100644 --- a/tank/model_utils.py +++ b/tank/model_utils.py @@ -1,10 +1,13 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers import torch import numpy as np import torchvision.models as models -from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel +from transformers import ( + AutoModelForSequenceClassification, + BertTokenizer, + TFBertModel, +) torch.manual_seed(0) @@ -12,17 +15,13 @@ torch.manual_seed(0) class HuggingFaceLanguage(torch.nn.Module): - def __init__(self, hf_model_name): super().__init__() self.model = AutoModelForSequenceClassification.from_pretrained( hf_model_name, # The pretrained model. - num_labels= - 2, # The number of output labels--2 for binary classification. - output_attentions= - False, # Whether the model returns attentions weights. - output_hidden_states= - False, # Whether the model returns all hidden-states. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. torchscript=True, ) @@ -44,7 +43,6 @@ def get_hf_model(name): class VisionModule(torch.nn.Module): - def __init__(self, model): super().__init__() self.model = model @@ -61,6 +59,7 @@ def get_vision_model(torch_model): actual_out = model(test_input) return model, test_input, actual_out + ################################################################################ # Utility function for comparing two tensors (torch). @@ -70,4 +69,3 @@ def compare_tensors(torch_tensor, numpy_tensor): atol = 1e-03 torch_to_numpy = torch_tensor.detach().numpy() return np.allclose(torch_to_numpy, numpy_tensor, rtol, atol) - diff --git a/tank/model_utils_tf.py b/tank/model_utils_tf.py index 46c04cd0..978cfc36 100644 --- a/tank/model_utils_tf.py +++ b/tank/model_utils_tf.py @@ -1,9 +1,12 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers import tensorflow as tf import numpy as np -from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel +from transformers import ( + AutoModelForSequenceClassification, + BertTokenizer, + TFBertModel, +) ##################### Tensorflow Hugging Face LM Models ################################### MAX_SEQUENCE_LENGTH = 512 @@ -13,20 +16,20 @@ BATCH_SIZE = 1 tf_bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] -class TFHuggingFaceLanguage(tf.Module): +class TFHuggingFaceLanguage(tf.Module): def __init__(self, hf_model_name): super(TFHuggingFaceLanguage, self).__init__() # Create a BERT trainer with the created network. - self.m = TFBertModel.from_pretrained( - hf_model_name, from_pt=True) + self.m = TFBertModel.from_pretrained(hf_model_name, from_pt=True) # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=tf_bert_input) def forward(self, input_ids, attention_mask, token_type_ids): @@ -36,17 +39,24 @@ class TFHuggingFaceLanguage(tf.Module): def get_TFhf_model(name): model = TFHuggingFaceLanguage(name) tokenizer = BertTokenizer.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased") + "microsoft/MiniLM-L12-H384-uncased" + ) text = "Replace me by any text you'd like." - encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + encoded_input = tokenizer( + text, + padding="max_length", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) - test_input = (encoded_input["input_ids"], encoded_input["attention_mask"], - encoded_input["token_type_ids"]) + tf.convert_to_tensor(encoded_input[key]), 0 + ) + test_input = ( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ) actual_out = model.forward(*test_input) return model, test_input, actual_out @@ -58,5 +68,3 @@ def compare_tensors_tf(tf_tensor, numpy_tensor): atol = 1e-03 tf_to_numpy = tf_tensor.numpy() return np.allclose(tf_to_numpy, numpy_tensor, rtol, atol) - - diff --git a/tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py b/tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py index 816b13f3..132f52c9 100644 --- a/tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py +++ b/tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_hf_model, compare_tensors from shark.parser import shark_args @@ -10,8 +10,8 @@ import pytest torch.manual_seed(0) -class MiniLMModuleTester: +class MiniLMModuleTester: def __init__( self, dynamic=False, @@ -25,24 +25,29 @@ class MiniLMModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_hf_model("microsoft/MiniLM-L12-H384-uncased") + model, input, act_out = get_hf_model( + "microsoft/MiniLM-L12-H384-uncased" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - shark_module = SharkInference(model, (input,), - device=self.device, - dynamic=self.dynamic, - jit_trace=True) + shark_module = SharkInference( + model, + (input,), + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class MiniLMModuleTest(unittest.TestCase): +class MiniLMModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") - + def setUp(self): self.module_tester = MiniLMModuleTester(self) @@ -50,20 +55,24 @@ class MiniLMModuleTest(unittest.TestCase): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + @pytest.mark.xfail(reason="language models failing for dynamic case") def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" @@ -71,9 +80,9 @@ class MiniLMModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" @@ -81,14 +90,14 @@ class MiniLMModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py b/tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py index 53044659..06f2ce3c 100644 --- a/tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py +++ b/tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_hf_model, compare_tensors from shark.parser import shark_args @@ -8,10 +8,10 @@ import unittest import numpy as np import pytest -#torch.manual_seed(0) +# torch.manual_seed(0) + class AlbertModuleTester: - def __init__( self, dynamic=False, @@ -28,21 +28,24 @@ class AlbertModuleTester: model, input, act_out = get_hf_model("albert-base-v2") shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - shark_module = SharkInference(model, (input,), - device=self.device, - dynamic=self.dynamic, - jit_trace=True) + shark_module = SharkInference( + model, + (input,), + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) + class AlbertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") - + def setUp(self): self.module_tester = AlbertModuleTester(self) self.module_tester.save_mlir = self.save_mlir @@ -51,20 +54,26 @@ class AlbertModuleTest(unittest.TestCase): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.xfail(reason="Language models currently failing for dynamic case") + + @pytest.mark.xfail( + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" @@ -72,9 +81,9 @@ class AlbertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" @@ -82,13 +91,14 @@ class AlbertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/alexnet/alexnet_pytorch_test.py b/tank/pytorch/alexnet/alexnet_pytorch_test.py index ca427908..1fbb3a0a 100644 --- a/tank/pytorch/alexnet/alexnet_pytorch_test.py +++ b/tank/pytorch/alexnet/alexnet_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class AlexnetModuleTester: +class AlexnetModuleTester: def __init__( self, dynamic=False, @@ -26,69 +26,75 @@ class AlexnetModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.alexnet(pretrained=True)) + model, input, act_out = get_vision_model( + models.alexnet(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class AlexnetModuleTest(unittest.TestCase): +class AlexnetModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") def setUp(self): self.module_tester = AlexnetModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py b/tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py index 927f3e05..2db9c6f1 100644 --- a/tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py +++ b/tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_hf_model, compare_tensors from shark.parser import shark_args @@ -8,10 +8,10 @@ import unittest import numpy as np import pytest -#torch.manual_seed(0) +# torch.manual_seed(0) + class BertModuleTester: - def __init__( self, dynamic=False, @@ -28,42 +28,51 @@ class BertModuleTester: model, input, act_out = get_hf_model("bert-base-uncased") shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - shark_module = SharkInference(model, (input,), - device=self.device, - dynamic=self.dynamic, - jit_trace=True) + shark_module = SharkInference( + model, + (input,), + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class BertModuleTest(unittest.TestCase): +class BertModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") - + def setUp(self): self.module_tester = BertModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.xfail(reason="Language models currently failing for dynamic case") + + @pytest.mark.xfail( + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" @@ -71,9 +80,9 @@ class BertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" @@ -81,14 +90,14 @@ class BertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py b/tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py index 6cc9fab1..9ba260f8 100644 --- a/tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py +++ b/tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_hf_model, compare_tensors from shark.parser import shark_args @@ -8,10 +8,10 @@ import unittest import numpy as np import pytest -#torch.manual_seed(0) +# torch.manual_seed(0) + class DistilBertModuleTester: - def __init__( self, dynamic=False, @@ -28,43 +28,54 @@ class DistilBertModuleTester: model, input, act_out = get_hf_model("distilbert-base-uncased") shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - shark_module = SharkInference(model, (input,), - device=self.device, - dynamic=self.dynamic, - jit_trace=True) + shark_module = SharkInference( + model, + (input,), + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class DistilBertModuleTest(unittest.TestCase): +class DistilBertModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") - + def setUp(self): self.module_tester = DistilBertModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.xfail(reason="Language models currently failing for dynamic case") + + @pytest.mark.xfail( + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") - @pytest.mark.xfail(reason="Language models currently failing for dynamic case") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) + @pytest.mark.xfail( + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" @@ -72,9 +83,9 @@ class DistilBertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" @@ -82,14 +93,14 @@ class DistilBertModuleTest(unittest.TestCase): @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554") @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/resnet101/resnet101_pytorch_test.py b/tank/pytorch/resnet101/resnet101_pytorch_test.py index 2a8c2448..af982650 100644 --- a/tank/pytorch/resnet101/resnet101_pytorch_test.py +++ b/tank/pytorch/resnet101/resnet101_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class Resnet101ModuleTester: +class Resnet101ModuleTester: def __init__( self, dynamic=False, @@ -26,68 +26,75 @@ class Resnet101ModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.resnet101(pretrained=True)) + model, input, act_out = get_vision_model( + models.resnet101(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class Resnet101ModuleTest(unittest.TestCase): +class Resnet101ModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") def setUp(self): self.module_tester = Resnet101ModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/resnet18/resnet18_pytorch_test.py b/tank/pytorch/resnet18/resnet18_pytorch_test.py index 4ce10516..a3372bee 100644 --- a/tank/pytorch/resnet18/resnet18_pytorch_test.py +++ b/tank/pytorch/resnet18/resnet18_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class Resnet18ModuleTester: +class Resnet18ModuleTester: def __init__( self, dynamic=False, @@ -26,69 +26,75 @@ class Resnet18ModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.resnet18(pretrained=True)) + model, input, act_out = get_vision_model( + models.resnet18(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class Resnet18ModuleTest(unittest.TestCase): +class Resnet18ModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") def setUp(self): self.module_tester = Resnet18ModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/resnet50/resnet50_pytorch_test.py b/tank/pytorch/resnet50/resnet50_pytorch_test.py index 6fbaf64b..4ac2e839 100644 --- a/tank/pytorch/resnet50/resnet50_pytorch_test.py +++ b/tank/pytorch/resnet50/resnet50_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class Resnet50ModuleTester: +class Resnet50ModuleTester: def __init__( self, dynamic=False, @@ -26,68 +26,75 @@ class Resnet50ModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.resnet50(pretrained=True)) + model, input, act_out = get_vision_model( + models.resnet50(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class Resnet50ModuleTest(unittest.TestCase): +class Resnet50ModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") + def setUp(self): self.module_tester = Resnet50ModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py b/tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py index fbd3f026..7d95b8ee 100644 --- a/tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py +++ b/tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class SqueezenetModuleTester: +class SqueezenetModuleTester: def __init__( self, dynamic=False, @@ -26,69 +26,75 @@ class SqueezenetModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.squeezenet1_0(pretrained=True)) + model, input, act_out = get_vision_model( + models.squeezenet1_0(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class SqueezenetModuleTest(unittest.TestCase): +class SqueezenetModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") def setUp(self): self.module_tester = SqueezenetModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/pytorch/v_diffusion/cfg_sample.py b/tank/pytorch/v_diffusion/cfg_sample.py index 80c221e0..71e0ef36 100644 --- a/tank/pytorch/v_diffusion/cfg_sample.py +++ b/tank/pytorch/v_diffusion/cfg_sample.py @@ -47,18 +47,23 @@ if os.path.exists(checkpoint): model.load_state_dict(torch.load(checkpoint, map_location="cpu")) model = model.to(device).eval().requires_grad_(False) -clip_model_name = model.clip_model if hasattr(model, "clip_model") else "ViT-B/16" +clip_model_name = ( + model.clip_model if hasattr(model, "clip_model") else "ViT-B/16" +) clip_model = clip.load(clip_model_name, jit=False, device=device)[0] clip_model.eval().requires_grad_(False) normalize = transforms.Normalize( - mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711] + mean=[0.48145466, 0.4578275, 0.40821073], + std=[0.26862954, 0.26130258, 0.27577711], ) zero_embed = torch.zeros([1, clip_model.visual.output_dim], device=device) target_embeds, weights = [zero_embed], [] txt, weight = parse_prompt(args.prompts[0]) -target_embeds.append(clip_model.encode_text(clip.tokenize(txt).to(device)).float()) +target_embeds.append( + clip_model.encode_text(clip.tokenize(txt).to(device)).float() +) weights.append(weight) weights = torch.tensor([1 - sum(weights), *weights], device=device) diff --git a/tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py b/tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py index 9be2b642..c40a59f3 100644 --- a/tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py +++ b/tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils import get_vision_model, compare_tensors from shark.parser import shark_args @@ -11,8 +11,8 @@ import pytest torch.manual_seed(0) -class WideResnet50ModuleTester: +class WideResnet50ModuleTester: def __init__( self, dynamic=False, @@ -26,69 +26,75 @@ class WideResnet50ModuleTester: self.save_vmfb = save_vmfb def create_and_check_module(self): - model, input, act_out = get_vision_model(models.wide_resnet50_2(pretrained=True)) + model, input, act_out = get_vision_model( + models.wide_resnet50_2(pretrained=True) + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb shark_module = SharkInference( - model, - (input,), - device=self.device, - dynamic=self.dynamic, + model, + (input,), + device=self.device, + dynamic=self.dynamic, ) shark_module.compile() results = shark_module.forward((input,)) assert True == compare_tensors(act_out, results) -class WideResnet50ModuleTest(unittest.TestCase): +class WideResnet50ModuleTest(unittest.TestCase): @pytest.fixture(autouse=True) - def configure(self, pytestconfig): + def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") self.save_vmfb = pytestconfig.getoption("save_vmfb") def setUp(self): self.module_tester = WideResnet50ModuleTester(self) - + def test_module_static_cpu(self): self.module_tester.dynamic = False self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - + def test_module_dynamic_cpu(self): self.module_tester.dynamic = True self.module_tester.device = "cpu" self.module_tester.create_and_check_module() - - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): self.module_tester.dynamic = False self.module_tester.device = "gpu" self.module_tester.create_and_check_module() - @pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): self.module_tester.dynamic = True self.module_tester.device = "gpu" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_static_vulkan(self): self.module_tester.dynamic = False self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() @pytest.mark.skipif( - check_device_drivers("vulkan"), - reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" - ) + check_device_drivers("vulkan"), + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", + ) def test_module_dynamic_vulkan(self): self.module_tester.dynamic = True self.module_tester.device = "vulkan" self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/automodelmaskedlm.py b/tank/tf/automodelmaskedlm.py index 18da1e0f..39a689b6 100644 --- a/tank/tf/automodelmaskedlm.py +++ b/tank/tf/automodelmaskedlm.py @@ -9,11 +9,11 @@ inputs_signature = [ class AutoModelMaskedLM(tf.Module): - def __init__(self, model_name): super(AutoModelMaskedLM, self).__init__() - self.m = TFAutoModelForMaskedLM.from_pretrained(model_name, - output_attentions=False) + self.m = TFAutoModelForMaskedLM.from_pretrained( + model_name, output_attentions=False + ) self.m.predict = lambda x: self.m(input_ids=x) @tf.function(input_signature=inputs_signature) @@ -24,20 +24,26 @@ class AutoModelMaskedLM(tf.Module): fail_models = ["microsoft/deberta-base", "google/rembert", "google/tapas-base"] supported_models = [ - "albert-base-v2", "bert-base-uncased", "camembert-base", - "dbmdz/convbert-base-turkish-cased", "distilbert-base-uncased", + "albert-base-v2", + "bert-base-uncased", + "camembert-base", + "dbmdz/convbert-base-turkish-cased", + "distilbert-base-uncased", "google/electra-small-discriminator", - "hf-internal-testing/tiny-random-flaubert", "funnel-transformer/small", - "microsoft/layoutlm-base-uncased", "allenai/longformer-base-4096", - "google/mobilebert-uncased", "microsoft/mpnet-base", "roberta-base", - "xlm-roberta-base" + "hf-internal-testing/tiny-random-flaubert", + "funnel-transformer/small", + "microsoft/layoutlm-base-uncased", + "allenai/longformer-base-4096", + "google/mobilebert-uncased", + "microsoft/mpnet-base", + "roberta-base", + "xlm-roberta-base", ] if __name__ == "__main__": - inputs = tf.random.uniform(shape=[1, 512], - maxval=3, - dtype=tf.int32, - seed=10) + inputs = tf.random.uniform( + shape=[1, 512], maxval=3, dtype=tf.int32, seed=10 + ) for model_name in supported_models: print(f"Running model: {model_name}") diff --git a/tank/tf/bert_large_gen.py b/tank/tf/bert_large_gen.py index d021daa3..e52379cf 100644 --- a/tank/tf/bert_large_gen.py +++ b/tank/tf/bert_large_gen.py @@ -19,24 +19,26 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=24, - hidden_size=1024, - num_attention_heads=16, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, + num_layers=24, + hidden_size=1024, + num_attention_heads=16, + dict_outputs=dict_outputs, + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -46,15 +48,20 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input0: input_word_ids - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input1: input_mask - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input2: segment_ids - tf.TensorSpec([BATCH_SIZE], tf.int32) # input3: labels - ]) + @tf.function( + input_signature=[ + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input0: input_word_ids + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input1: input_mask + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input2: segment_ids + tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels + ] + ) def learn(self, input_word_ids, input_mask, segment_ids, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -77,12 +84,12 @@ class BertModule(tf.Module): if __name__ == "__main__": # BertModule() # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["learn"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["learn"], import_only=True + ) # Save module as MLIR file in a directory ARITFACTS_DIR = os.getcwd() mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir") with open(mlir_path, "wt") as output_file: - output_file.write(compiler_module.decode('utf-8')) + output_file.write(compiler_module.decode("utf-8")) print(f"Wrote MLIR to path '{mlir_path}'") diff --git a/tank/tf/bert_large_run.py b/tank/tf/bert_large_run.py index dc5cf07e..79c152bb 100644 --- a/tank/tf/bert_large_run.py +++ b/tank/tf/bert_large_run.py @@ -21,24 +21,26 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=24, - hidden_size=1024, - num_attention_heads=16, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, + num_layers=24, + hidden_size=1024, + num_attention_heads=16, + dict_outputs=dict_outputs, + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -49,10 +51,12 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - bert_input, # inputs - tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels - ]) + @tf.function( + input_signature=[ + bert_input, # inputs + tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels + ] + ) def learn(self, inputs, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -69,26 +73,29 @@ class BertModule(tf.Module): if __name__ == "__main__": # BertModule() # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["learn"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["learn"], import_only=True + ) # Compile the model using IREE backend = "dylib-llvm-aot" args = [ "--iree-llvm-target-cpu-features=host", "--iree-mhlo-demote-i64-to-i32=false", - "--iree-stream-resource-index-bits=64", "--iree-vm-target-index-bits=64" + "--iree-stream-resource-index-bits=64", + "--iree-vm-target-index-bits=64", ] backend_config = "dylib" - #backend = "cuda" - #backend_config = "cuda" - #args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] - flatbuffer_blob = compile_str(compiler_module, - target_backends=[backend], - extra_args=args, - input_type="mhlo") - #flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) + # backend = "cuda" + # backend_config = "cuda" + # args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] + flatbuffer_blob = compile_str( + compiler_module, + target_backends=[backend], + extra_args=args, + input_type="mhlo", + ) + # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) # Save module as MLIR file in a directory vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob) @@ -100,21 +107,23 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), ] learn_sample_input = [ predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)) + np.random.randint(5, size=(BATCH_SIZE)), ] warmup = 5 total_iter = 10 num_iter = total_iter - warmup for i in range(10): - if (i == warmup - 1): + if i == warmup - 1: start = time.time() print( - BertCompiled.learn(predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)))) + BertCompiled.learn( + predict_sample_input, np.random.randint(5, size=(BATCH_SIZE)) + ) + ) end = time.time() total_time = end - start print("time: " + str(total_time)) diff --git a/tank/tf/bert_large_tf.py b/tank/tf/bert_large_tf.py index 09555502..4d35cfd8 100644 --- a/tank/tf/bert_large_tf.py +++ b/tank/tf/bert_large_tf.py @@ -14,24 +14,26 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=24, - hidden_size=1024, - num_attention_heads=16, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, + num_layers=24, + hidden_size=1024, + num_attention_heads=16, + dict_outputs=dict_outputs, + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -42,10 +44,12 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - bert_input, # inputs - tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels - ]) + @tf.function( + input_signature=[ + bert_input, # inputs + tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels + ] + ) def learn(self, inputs, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -64,7 +68,7 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), ] bert_model = BertModule() warmup = 1 @@ -72,9 +76,11 @@ if __name__ == "__main__": num_iter = total_iter - warmup for i in range(total_iter): print( - bert_model.learn(predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)))) - if (i == warmup - 1): + bert_model.learn( + predict_sample_input, np.random.randint(5, size=(BATCH_SIZE)) + ) + ) + if i == warmup - 1: start = time.time() end = time.time() diff --git a/tank/tf/bert_small_gen.py b/tank/tf/bert_small_gen.py index d56ca304..d39af7d2 100644 --- a/tank/tf/bert_small_gen.py +++ b/tank/tf/bert_small_gen.py @@ -1,5 +1,6 @@ from iree import runtime as ireert -#from iree.tf.support import module_utils + +# from iree.tf.support import module_utils from iree.compiler import tf as tfc from absl import app @@ -19,22 +20,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=2, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -44,15 +45,20 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input0: input_word_ids - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input1: input_mask - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], - dtype=tf.int32), #input2: segment_ids - tf.TensorSpec([BATCH_SIZE], tf.int32) # input3: labels - ]) + @tf.function( + input_signature=[ + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input0: input_word_ids + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input1: input_mask + tf.TensorSpec( + shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32 + ), # input2: segment_ids + tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels + ] + ) def learn(self, input_word_ids, input_mask, segment_ids, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -75,13 +81,13 @@ class BertModule(tf.Module): if __name__ == "__main__": # BertModule() # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["learn"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["learn"], import_only=True + ) print(type(compiler_module)) # Save module as MLIR file in a directory ARITFACTS_DIR = os.getcwd() mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir") with open(mlir_path, "wt") as output_file: - output_file.write(compiler_module.decode('utf-8')) + output_file.write(compiler_module.decode("utf-8")) print(f"Wrote MLIR to path '{mlir_path}'") diff --git a/tank/tf/bert_small_run.py b/tank/tf/bert_small_run.py index 90f8375e..47af7dae 100644 --- a/tank/tf/bert_small_run.py +++ b/tank/tf/bert_small_run.py @@ -21,22 +21,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=2, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -47,10 +47,12 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - bert_input, # inputs - tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels - ]) + @tf.function( + input_signature=[ + bert_input, # inputs + tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels + ] + ) def learn(self, inputs, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -67,25 +69,28 @@ class BertModule(tf.Module): if __name__ == "__main__": # BertModule() # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["learn"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["learn"], import_only=True + ) # Compile the model using IREE backend = "dylib-llvm-aot" args = [ "--iree-llvm-target-cpu-features=host", - "--iree-mhlo-demote-i64-to-i32=false", "--iree-flow-demote-i64-to-i32" + "--iree-mhlo-demote-i64-to-i32=false", + "--iree-flow-demote-i64-to-i32", ] backend_config = "dylib" - #backend = "cuda" - #backend_config = "cuda" - #args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] - flatbuffer_blob = compile_str(compiler_module, - target_backends=[backend], - extra_args=args, - input_type="mhlo") - #flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) + # backend = "cuda" + # backend_config = "cuda" + # args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] + flatbuffer_blob = compile_str( + compiler_module, + target_backends=[backend], + extra_args=args, + input_type="mhlo", + ) + # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) # Save module as MLIR file in a directory vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob) @@ -97,21 +102,23 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), ] learn_sample_input = [ predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)) + np.random.randint(5, size=(BATCH_SIZE)), ] warmup = 5 total_iter = 10 num_iter = total_iter - warmup for i in range(10): - if (i == warmup - 1): + if i == warmup - 1: start = time.time() print( - BertCompiled.learn(predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)))) + BertCompiled.learn( + predict_sample_input, np.random.randint(5, size=(BATCH_SIZE)) + ) + ) end = time.time() total_time = end - start print("time: " + str(total_time)) diff --git a/tank/tf/bert_small_tf_run.py b/tank/tf/bert_small_tf_run.py index add3cb40..18f665b8 100644 --- a/tank/tf/bert_small_tf_run.py +++ b/tank/tf/bert_small_tf_run.py @@ -14,22 +14,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() dict_outputs = False - test_network = networks.BertEncoder(vocab_size=vocab_size, - num_layers=2, - dict_outputs=dict_outputs) + test_network = networks.BertEncoder( + vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs + ) # Create a BERT trainer with the created network. bert_trainer_model = bert_classifier.BertClassifier( - test_network, num_classes=NUM_CLASSES) + test_network, num_classes=NUM_CLASSES + ) bert_trainer_model.summary() # Invoke the trainer model on the inputs. This causes the layer to be built. @@ -40,10 +40,12 @@ class BertModule(tf.Module): self.loss = tf.keras.losses.SparseCategoricalCrossentropy() self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2) - @tf.function(input_signature=[ - bert_input, # inputs - tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels - ]) + @tf.function( + input_signature=[ + bert_input, # inputs + tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels + ] + ) def learn(self, inputs, labels): with tf.GradientTape() as tape: # Capture the gradients from forward prop... @@ -62,7 +64,7 @@ if __name__ == "__main__": predict_sample_input = [ np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), - np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)) + np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)), ] bert_model = BertModule() warmup = 1 @@ -70,9 +72,11 @@ if __name__ == "__main__": num_iter = total_iter - warmup for i in range(total_iter): print( - bert_model.learn(predict_sample_input, - np.random.randint(5, size=(BATCH_SIZE)))) - if (i == warmup - 1): + bert_model.learn( + predict_sample_input, np.random.randint(5, size=(BATCH_SIZE)) + ) + ) + if i == warmup - 1: start = time.time() end = time.time() diff --git a/tank/tf/hf_masked_lm/albert-base-v2_test.py b/tank/tf/hf_masked_lm/albert-base-v2_test.py index 901d8fbd..d10f6384 100644 --- a/tank/tf/hf_masked_lm/albert-base-v2_test.py +++ b/tank/tf/hf_masked_lm/albert-base-v2_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class AlbertBaseModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -42,71 +41,80 @@ class AlbertBaseModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "albert-base-v2", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "albert-base-v2", dynamic, device + ) class AlbertBaseModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = AlbertBaseModuleTester(self) - self.module_tester.save_temps=pytestconfig.getoption("save_temps") - self.module_tester.save_mlir=pytestconfig.getoption("save_mlir") - self.module_tester.save_vmfb=pytestconfig.getoption("save_vmfb") - self.module_tester.benchmark=pytestconfig.getoption("benchmark") + self.module_tester.save_temps = pytestconfig.getoption("save_temps") + self.module_tester.save_mlir = pytestconfig.getoption("save_mlir") + self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") + self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) - + @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.xfail( - reason="https://github.com/google/iree/issues/9553") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.xfail(reason="https://github.com/google/iree/issues/9553") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -115,8 +123,7 @@ class AlbertBaseModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -124,11 +131,11 @@ class AlbertBaseModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -136,5 +143,5 @@ class AlbertBaseModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/bert-base-uncased_test.py b/tank/tf/hf_masked_lm/bert-base-uncased_test.py index 24c6a25b..2650fd87 100644 --- a/tank/tf/hf_masked_lm/bert-base-uncased_test.py +++ b/tank/tf/hf_masked_lm/bert-base-uncased_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class BertBaseUncasedModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -42,36 +41,40 @@ class BertBaseUncasedModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "bert_base_uncased", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "bert_base_uncased", dynamic, device + ) class BertBaseUncasedModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = BertBaseUncasedModuleTester(self) @@ -80,8 +83,9 @@ class BertBaseUncasedModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" @@ -89,24 +93,28 @@ class BertBaseUncasedModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -115,8 +123,7 @@ class BertBaseUncasedModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -124,11 +131,11 @@ class BertBaseUncasedModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -136,5 +143,5 @@ class BertBaseUncasedModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/camembert-base_test.py b/tank/tf/hf_masked_lm/camembert-base_test.py index 8bffff36..bf48b959 100644 --- a/tank/tf/hf_masked_lm/camembert-base_test.py +++ b/tank/tf/hf_masked_lm/camembert-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class CamemBertModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -42,44 +41,51 @@ class CamemBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - + if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "camembert-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "camembert-base", dynamic, device + ) + class CamemBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): - self.module_tester=CamemBertModuleTester(self) + self.module_tester = CamemBertModuleTester(self) self.module_tester.save_temps = pytestconfig.getoption("save_temps") self.module_tester.save_vmfb = pytestconfig.getoption("save_mlir") self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" @@ -87,24 +93,28 @@ class CamemBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -113,8 +123,7 @@ class CamemBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -122,11 +131,11 @@ class CamemBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -134,5 +143,5 @@ class CamemBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py b/tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py index 7f203ec7..aac8f5f6 100644 --- a/tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py +++ b/tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -11,24 +11,23 @@ import numpy as np import tempfile -class ConvBertModuleTester: - +class ConvBertModuleTester: def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb self.benchmark = benchmark - def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model( - "dbmdz/convbert-base-turkish-cased") + "dbmdz/convbert-base-turkish-cased" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb if self.save_temps == True: @@ -44,36 +43,40 @@ class ConvBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "convbert-base-turkish-cased", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "convbert-base-turkish-cased", dynamic, device + ) class ConvBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = ConvBertModuleTester(self) @@ -82,7 +85,9 @@ class ConvBertModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536.") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536." + ) def test_module_static_cpu(self): dynamic = False device = "cpu" @@ -90,24 +95,28 @@ class ConvBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -116,8 +125,7 @@ class ConvBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -125,11 +133,11 @@ class ConvBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -137,5 +145,5 @@ class ConvBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/deberta-base_test.py b/tank/tf/hf_masked_lm/deberta-base_test.py index 410d6d86..29c96d32 100644 --- a/tank/tf/hf_masked_lm/deberta-base_test.py +++ b/tank/tf/hf_masked_lm/deberta-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class DebertaModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("microsoft/deberta-base") shark_args.save_mlir = self.save_mlir @@ -41,36 +40,40 @@ class DebertaModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "deberta-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "deberta-base", dynamic, device + ) class DebertaModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = DebertaModuleTester(self) @@ -78,10 +81,11 @@ class DebertaModuleTest(unittest.TestCase): self.module_tester.save_mlir = pytestconfig.getoption("save_mlir") self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - + @pytest.mark.xfail - @pytest.mark.skip(reason="deberta currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="deberta currently failing in the lowering passes." + ) def test_module_static_cpu(self): dynamic = False device = "cpu" @@ -89,52 +93,59 @@ class DebertaModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="deberta currently failing in the lowering passes." - ) - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skip( + reason="deberta currently failing in the lowering passes." + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="deberta currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="deberta currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="deberta currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="deberta currently failing in the lowering passes." + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False device = "vulkan" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="deberta currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="deberta currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -142,5 +153,5 @@ class DebertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/distilbert-base-uncased_test.py b/tank/tf/hf_masked_lm/distilbert-base-uncased_test.py index c4d0e06f..b4ef0f9b 100644 --- a/tank/tf/hf_masked_lm/distilbert-base-uncased_test.py +++ b/tank/tf/hf_masked_lm/distilbert-base-uncased_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class DistilBertModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -41,35 +40,40 @@ class DistilBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "distilbert-base-uncased", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "distilbert-base-uncased", dynamic, device + ) + class DistilBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = DistilBertModuleTester(self) @@ -78,7 +82,9 @@ class DistilBertModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" @@ -86,24 +92,28 @@ class DistilBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -112,8 +122,7 @@ class DistilBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -121,11 +130,11 @@ class DistilBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -133,5 +142,5 @@ class DistilBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/electra-small-discriminator_test.py b/tank/tf/hf_masked_lm/electra-small-discriminator_test.py index 66b77168..febcc9d5 100644 --- a/tank/tf/hf_masked_lm/electra-small-discriminator_test.py +++ b/tank/tf/hf_masked_lm/electra-small-discriminator_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,20 +12,21 @@ import tempfile class ElectraModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb def create_and_check_module(self, dynamic, device): - model, input, act_out = get_causal_lm_model("google/electra-small-discriminator") + model, input, act_out = get_causal_lm_model( + "google/electra-small-discriminator" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb if self.save_temps == True: @@ -39,71 +40,81 @@ class ElectraModuleTester: np.save(f"{temp_dir}/input2.npy", input[1]) exp_out = act_out.numpy() with open(f"{temp_dir}/expected_out.txt", "w") as out_file: - out_file.write(np.array2string(exp_out)) + out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - + if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "electra-small-discriminator", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "electra-small-discriminator", dynamic, device + ) + class ElectraModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = ElectraModuleTester(self) - self.module_tester.save_temps = pytestconfig.getoption("save_temps") + self.module_tester.save_temps = pytestconfig.getoption("save_temps") self.module_tester.save_mlir = pytestconfig.getoption("save_mlir") self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -112,8 +123,7 @@ class ElectraModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -121,11 +131,11 @@ class ElectraModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -133,5 +143,5 @@ class ElectraModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/funnel-transformer_test.py b/tank/tf/hf_masked_lm/funnel-transformer_test.py index 7c60ed58..47c79f58 100644 --- a/tank/tf/hf_masked_lm/funnel-transformer_test.py +++ b/tank/tf/hf_masked_lm/funnel-transformer_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class FunnelModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -41,35 +40,40 @@ class FunnelModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "funnel-transformer-small", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "funnel-transformer-small", dynamic, device + ) + class FunnelModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = FunnelModuleTester(self) @@ -85,15 +89,17 @@ class FunnelModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip(reason="funnel currently failing in the lowering passes.") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" @@ -101,9 +107,11 @@ class FunnelModuleTest(unittest.TestCase): @pytest.mark.skip(reason="funnel currently failing in the lowering passes.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -112,8 +120,7 @@ class FunnelModuleTest(unittest.TestCase): @pytest.mark.skip(reason="funnel currently failing in the lowering passes.") @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -122,11 +129,11 @@ class FunnelModuleTest(unittest.TestCase): @pytest.mark.skip(reason="funnel currently failing in the lowering passes.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -134,5 +141,5 @@ class FunnelModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/layoutlm-base-uncased_test.py b/tank/tf/hf_masked_lm/layoutlm-base-uncased_test.py index b6bbfcf7..9b128b3f 100644 --- a/tank/tf/hf_masked_lm/layoutlm-base-uncased_test.py +++ b/tank/tf/hf_masked_lm/layoutlm-base-uncased_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,21 +12,21 @@ import tempfile class LayoutLmModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model( - "microsoft/layoutlm-base-uncased") + "microsoft/layoutlm-base-uncased" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb if self.save_temps == True: @@ -42,35 +42,40 @@ class LayoutLmModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "layoutlm-base-uncased", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "layoutlm-base-uncased", dynamic, device + ) + class LayoutLmModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = LayoutLmModuleTester(self) @@ -78,32 +83,38 @@ class LayoutLmModuleTest(unittest.TestCase): self.module_tester.save_mlir = pytestconfig.getoption("save_mlir") self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -112,8 +123,7 @@ class LayoutLmModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -121,11 +131,11 @@ class LayoutLmModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -133,5 +143,5 @@ class LayoutLmModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/longformer-base-4096_test.py b/tank/tf/hf_masked_lm/longformer-base-4096_test.py index 4fac8fd4..b1ccc132 100644 --- a/tank/tf/hf_masked_lm/longformer-base-4096_test.py +++ b/tank/tf/hf_masked_lm/longformer-base-4096_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,20 +12,21 @@ import tempfile class LongFormerModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): - model, input, act_out = get_causal_lm_model("allenai/longformer-base-4096") + model, input, act_out = get_causal_lm_model( + "allenai/longformer-base-4096" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb if self.save_temps == True: @@ -41,35 +42,40 @@ class LongFormerModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "longformer-base-4096", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "longformer-base-4096", dynamic, device + ) + class LongFormerModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = LongFormerModuleTester(self) @@ -79,45 +85,52 @@ class LongFormerModuleTest(unittest.TestCase): self.module_tester.benchmark = pytestconfig.getoption("benchmark") @pytest.mark.skip( - reason="longformer currently failing in the lowering passes.") + reason="longformer currently failing in the lowering passes." + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="longformer currently failing in the lowering passes.") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="longformer currently failing in the lowering passes." + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="longformer currently failing in the lowering passes.") + reason="longformer currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="longformer currently failing in the lowering passes.") + reason="longformer currently failing in the lowering passes." + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -125,13 +138,14 @@ class LongFormerModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="longformer currently failing in the lowering passes.") + reason="longformer currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -139,5 +153,5 @@ class LongFormerModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/masked_lm.py b/tank/tf/hf_masked_lm/masked_lm.py index 6dac73d3..4e309719 100644 --- a/tank/tf/hf_masked_lm/masked_lm.py +++ b/tank/tf/hf_masked_lm/masked_lm.py @@ -17,21 +17,22 @@ inputs_signature = [ def preprocess_input(model_name, text="This is just used to compile the model"): tokenizer = AutoTokenizer.from_pretrained(model_name) - inputs = tokenizer(text, - padding="max_length", - return_tensors="tf", - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + inputs = tokenizer( + text, + padding="max_length", + return_tensors="tf", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) return inputs class MaskedLM(tf.Module): - def __init__(self, model_name): super(MaskedLM, self).__init__() - self.m = TFAutoModelForMaskedLM.from_pretrained(model_name, - output_attentions=False, - num_labels=2) + self.m = TFAutoModelForMaskedLM.from_pretrained( + model_name, output_attentions=False, num_labels=2 + ) self.m.predict = lambda x, y: self.m(input_ids=x, attention_mask=y)[0] @tf.function(input_signature=inputs_signature) diff --git a/tank/tf/hf_masked_lm/mobilebert-uncased_test.py b/tank/tf/hf_masked_lm/mobilebert-uncased_test.py index aafb86df..c2b22dd7 100644 --- a/tank/tf/hf_masked_lm/mobilebert-uncased_test.py +++ b/tank/tf/hf_masked_lm/mobilebert-uncased_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,13 +12,12 @@ import tempfile class MobileBertModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir @@ -41,35 +40,40 @@ class MobileBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "mobilebert-uncased", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "mobilebert-uncased", dynamic, device + ) + class MobileBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = MobileBertModuleTester(self) @@ -78,31 +82,37 @@ class MobileBertModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -111,8 +121,7 @@ class MobileBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -120,11 +129,11 @@ class MobileBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -132,5 +141,5 @@ class MobileBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/mpnet-base_test.py b/tank/tf/hf_masked_lm/mpnet-base_test.py index eb5895bf..7cdf778b 100644 --- a/tank/tf/hf_masked_lm/mpnet-base_test.py +++ b/tank/tf/hf_masked_lm/mpnet-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class MpNetModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("microsoft/mpnet-base") shark_args.save_mlir = self.save_mlir @@ -41,35 +40,40 @@ class MpNetModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "mpnet-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "mpnet-base", dynamic, device + ) + class MpNetModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = MpNetModuleTester(self) @@ -78,31 +82,37 @@ class MpNetModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -111,8 +121,7 @@ class MpNetModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -120,11 +129,11 @@ class MpNetModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -132,5 +141,5 @@ class MpNetModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/rembert_test.py b/tank/tf/hf_masked_lm/rembert_test.py index 0d097b71..0e55a38b 100644 --- a/tank/tf/hf_masked_lm/rembert_test.py +++ b/tank/tf/hf_masked_lm/rembert_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class RemBertModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("google/rembert") shark_args.save_mlir = self.save_mlir @@ -41,35 +40,38 @@ class RemBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "rembert", - dynamic, - device) + shark_module.benchmark_all_csv((input), "rembert", dynamic, device) + class RemBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = RemBertModuleTester(self) @@ -78,60 +80,68 @@ class RemBertModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.skip(reason="rembert currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="rembert currently failing in the lowering passes." + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="rembert currently failing in the lowering passes." - ) - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skip( + reason="rembert currently failing in the lowering passes." + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="rembert currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="rembert currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="rembert currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="rembert currently failing in the lowering passes." + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False device = "vulkan" self.module_tester.create_and_check_module(dynamic, device) - @pytest.mark.skip(reason="rembert currently failing in the lowering passes." - ) + @pytest.mark.skip( + reason="rembert currently failing in the lowering passes." + ) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -139,5 +149,5 @@ class RemBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/roberta-base_test.py b/tank/tf/hf_masked_lm/roberta-base_test.py index 1dfc1cc8..1d5d0e47 100644 --- a/tank/tf/hf_masked_lm/roberta-base_test.py +++ b/tank/tf/hf_masked_lm/roberta-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class RobertaModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("roberta-base") shark_args.save_mlir = self.save_mlir @@ -41,35 +40,40 @@ class RobertaModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "roberta-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "roberta-base", dynamic, device + ) + class RobertaModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = RobertaModuleTester(self) @@ -78,31 +82,37 @@ class RobertaModuleTest(unittest.TestCase): self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb") self.module_tester.benchmark = pytestconfig.getoption("benchmark") - @pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536") + @pytest.mark.xfail( + reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536" + ) def test_module_static_cpu(self): dynamic = False device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -111,8 +121,7 @@ class RobertaModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -120,11 +129,11 @@ class RobertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -132,5 +141,5 @@ class RobertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/tapas-base_test.py b/tank/tf/hf_masked_lm/tapas-base_test.py index 4c2473ca..029495e3 100644 --- a/tank/tf/hf_masked_lm/tapas-base_test.py +++ b/tank/tf/hf_masked_lm/tapas-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class TapasBaseModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("google/tapas-base") shark_args.save_mlir = self.save_mlir @@ -41,35 +40,40 @@ class TapasBaseModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "tapas-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "tapas-base", dynamic, device + ) + class TapasBaseModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = TapasBaseModuleTester(self) @@ -85,15 +89,17 @@ class TapasBaseModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip(reason="tapas currently failing in the lowering passes.") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" @@ -101,9 +107,11 @@ class TapasBaseModuleTest(unittest.TestCase): @pytest.mark.skip(reason="tapas currently failing in the lowering passes.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -112,8 +120,7 @@ class TapasBaseModuleTest(unittest.TestCase): @pytest.mark.skip(reason="tapas currently failing in the lowering passes.") @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -122,11 +129,11 @@ class TapasBaseModuleTest(unittest.TestCase): @pytest.mark.skip(reason="tapas currently failing in the lowering passes.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -134,5 +141,5 @@ class TapasBaseModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/tiny-random-flaubert_test.py b/tank/tf/hf_masked_lm/tiny-random-flaubert_test.py index d5da7bfd..049893fb 100644 --- a/tank/tf/hf_masked_lm/tiny-random-flaubert_test.py +++ b/tank/tf/hf_masked_lm/tiny-random-flaubert_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,20 +12,21 @@ import tempfile class FlauBertModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): - model, input, act_out = get_causal_lm_model("hf-internal-testing/tiny-random-flaubert") + model, input, act_out = get_causal_lm_model( + "hf-internal-testing/tiny-random-flaubert" + ) shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb if self.save_temps == True: @@ -41,35 +42,40 @@ class FlauBertModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "tiny-random-flaubert", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "tiny-random-flaubert", dynamic, device + ) + class FlauBertModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = FlauBertModuleTester(self) @@ -84,24 +90,28 @@ class FlauBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -110,8 +120,7 @@ class FlauBertModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -119,11 +128,11 @@ class FlauBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -131,5 +140,5 @@ class FlauBertModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/hf_masked_lm/xlm-roberta-base_test.py b/tank/tf/hf_masked_lm/xlm-roberta-base_test.py index ff44ecc0..48fef0b5 100644 --- a/tank/tf/hf_masked_lm/xlm-roberta-base_test.py +++ b/tank/tf/hf_masked_lm/xlm-roberta-base_test.py @@ -1,6 +1,6 @@ from masked_lm import get_causal_lm_model from tank.model_utils_tf import compare_tensors_tf -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from shark.shark_inference import SharkInference from shark.parser import shark_args @@ -12,18 +12,17 @@ import tempfile class XLMRobertaModuleTester: - def __init__( self, save_temps=False, save_mlir=False, save_vmfb=False, - benchmark=False + benchmark=False, ): self.save_temps = save_temps self.save_mlir = save_mlir self.save_vmfb = save_vmfb - + def create_and_check_module(self, dynamic, device): model, input, act_out = get_causal_lm_model("xlm-roberta-base") shark_args.save_mlir = self.save_mlir @@ -41,35 +40,40 @@ class XLMRobertaModuleTester: with open(f"{temp_dir}/expected_out.txt", "w") as out_file: out_file.write(np.array2string(exp_out)) with ireec.tools.TempFileSaver(temp_dir): - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) - - else: - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True, - benchmark_mode=self.benchmark) + + else: + shark_module = SharkInference( + model, + (input,), + device=device, + dynamic=dynamic, + jit_trace=True, + benchmark_mode=self.benchmark, + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) assert True == compare_tensors_tf(act_out, results) if self.benchmark == True: - shark_module.benchmark_all_csv((input), - "xlm-roberta-base", - dynamic, - device) + shark_module.benchmark_all_csv( + (input), "xlm-roberta-base", dynamic, device + ) + class XLMRobertaModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.module_tester = XLMRobertaModuleTester(self) @@ -85,24 +89,28 @@ class XLMRobertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -111,8 +119,7 @@ class XLMRobertaModuleTest(unittest.TestCase): @pytest.mark.xfail @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -120,11 +127,11 @@ class XLMRobertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -132,5 +139,5 @@ class XLMRobertaModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/huggingface_MiniLM_gen.py b/tank/tf/huggingface_MiniLM_gen.py index 221aa0a0..46a3aa19 100644 --- a/tank/tf/huggingface_MiniLM_gen.py +++ b/tank/tf/huggingface_MiniLM_gen.py @@ -15,21 +15,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() # Create a BERT trainer with the created network. self.m = TFBertModel.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased", from_pt=True) + "microsoft/MiniLM-L12-H384-uncased", from_pt=True + ) # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=bert_input) def predict(self, input_word_ids, input_mask, segment_ids): @@ -39,12 +40,12 @@ class BertModule(tf.Module): if __name__ == "__main__": # BertModule() # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["predict"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["predict"], import_only=True + ) # Save module as MLIR file in a directory ARITFACTS_DIR = os.getcwd() mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir") with open(mlir_path, "wt") as output_file: - output_file.write(compiler_module.decode('utf-8')) + output_file.write(compiler_module.decode("utf-8")) print(f"Wrote MLIR to path '{mlir_path}'") diff --git a/tank/tf/huggingface_MiniLM_run.py b/tank/tf/huggingface_MiniLM_run.py index d01e4ec1..e989bc72 100644 --- a/tank/tf/huggingface_MiniLM_run.py +++ b/tank/tf/huggingface_MiniLM_run.py @@ -16,21 +16,22 @@ BATCH_SIZE = 1 bert_input = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), - tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32) + tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] class BertModule(tf.Module): - def __init__(self): super(BertModule, self).__init__() # Create a BERT trainer with the created network. self.m = TFBertModel.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased", from_pt=True) + "microsoft/MiniLM-L12-H384-uncased", from_pt=True + ) # Invoke the trainer model on the inputs. This causes the layer to be built. self.m.predict = lambda x, y, z: self.m.call( - input_ids=x, attention_mask=y, token_type_ids=z, training=False) + input_ids=x, attention_mask=y, token_type_ids=z, training=False + ) @tf.function(input_signature=bert_input) def predict(self, input_ids, attention_mask, token_type_ids): @@ -40,36 +41,43 @@ class BertModule(tf.Module): if __name__ == "__main__": # Prepping Data tokenizer = BertTokenizer.from_pretrained( - "microsoft/MiniLM-L12-H384-uncased") + "microsoft/MiniLM-L12-H384-uncased" + ) text = "Replace me by any text you'd like." - encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + encoded_input = tokenizer( + text, + padding="max_length", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) + tf.convert_to_tensor(encoded_input[key]), 0 + ) # Compile the model using IREE - compiler_module = tfc.compile_module(BertModule(), - exported_names=["predict"], - import_only=True) + compiler_module = tfc.compile_module( + BertModule(), exported_names=["predict"], import_only=True + ) # Compile the model using IREE backend = "dylib-llvm-aot" args = [ "--iree-llvm-target-cpu-features=host", - "--iree-mhlo-demote-i64-to-i32=false", "--iree-flow-demote-i64-to-i32" + "--iree-mhlo-demote-i64-to-i32=false", + "--iree-flow-demote-i64-to-i32", ] backend_config = "dylib" - #backend = "cuda" - #backend_config = "cuda" - #args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] - flatbuffer_blob = compile_str(compiler_module, - target_backends=[backend], - extra_args=args, - input_type="mhlo") - #flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) + # backend = "cuda" + # backend_config = "cuda" + # args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"] + flatbuffer_blob = compile_str( + compiler_module, + target_backends=[backend], + extra_args=args, + input_type="mhlo", + ) + # flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"]) # Save module as MLIR file in a directory vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob) @@ -78,7 +86,9 @@ if __name__ == "__main__": ctx = ireert.SystemContext(config=config) ctx.add_vm_module(vm_module) BertCompiled = ctx.modules.module - result = BertCompiled.predict(encoded_input["input_ids"], - encoded_input["attention_mask"], - encoded_input["token_type_ids"]) + result = BertCompiled.predict( + encoded_input["input_ids"], + encoded_input["attention_mask"], + encoded_input["token_type_ids"], + ) print(result) diff --git a/tank/tf/huggingface_MiniLM_tf.py b/tank/tf/huggingface_MiniLM_tf.py index a7204378..67eb4cb4 100644 --- a/tank/tf/huggingface_MiniLM_tf.py +++ b/tank/tf/huggingface_MiniLM_tf.py @@ -1,18 +1,19 @@ import tensorflow as tf from transformers import BertModel, BertTokenizer, TFBertModel -tf_model = TFBertModel.from_pretrained("microsoft/MiniLM-L12-H384-uncased", - from_pt=True) +tf_model = TFBertModel.from_pretrained( + "microsoft/MiniLM-L12-H384-uncased", from_pt=True +) tokenizer = BertTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") text = "Replace me by any text you'd like." -encoded_input = tokenizer(text, - padding='max_length', - truncation=True, - max_length=512) +encoded_input = tokenizer( + text, padding="max_length", truncation=True, max_length=512 +) for key in encoded_input: encoded_input[key] = tf.expand_dims( - tf.convert_to_tensor(encoded_input[key]), 0) + tf.convert_to_tensor(encoded_input[key]), 0 + ) output = tf_model(encoded_input) print(output) diff --git a/tank/tf/minilm_tf_test.py b/tank/tf/minilm_tf_test.py index 683aa941..10cc50f9 100644 --- a/tank/tf/minilm_tf_test.py +++ b/tank/tf/minilm_tf_test.py @@ -1,5 +1,5 @@ from shark.shark_inference import SharkInference -from shark.iree_utils import check_device_drivers +from shark.iree_utils._common import check_device_drivers from tank.model_utils_tf import get_TFhf_model, compare_tensors_tf import tensorflow as tf @@ -9,14 +9,13 @@ import pytest class MiniLMTFModuleTester: - def create_and_check_module(self, dynamic, device): model, input, act_out = get_TFhf_model( - "microsoft/MiniLM-L12-H384-uncased") - shark_module = SharkInference(model, (input,), - device=device, - dynamic=dynamic, - jit_trace=True) + "microsoft/MiniLM-L12-H384-uncased" + ) + shark_module = SharkInference( + model, (input,), device=device, dynamic=dynamic, jit_trace=True + ) shark_module.set_frontend("tensorflow") shark_module.compile() results = shark_module.forward((input)) @@ -24,7 +23,6 @@ class MiniLMTFModuleTester: class MiniLMTFModuleTest(unittest.TestCase): - def setUp(self): self.module_tester = MiniLMTFModuleTester() @@ -36,15 +34,17 @@ class MiniLMTFModuleTest(unittest.TestCase): @pytest.mark.skip(reason="TF testing temporarily unavailable.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) def test_module_dynamic_cpu(self): dynamic = True device = "cpu" self.module_tester.create_and_check_module(dynamic, device) @pytest.mark.skip(reason="TF testing temporarily unavailable.") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_static_gpu(self): dynamic = False device = "gpu" @@ -52,9 +52,11 @@ class MiniLMTFModuleTest(unittest.TestCase): @pytest.mark.skip(reason="TF testing temporarily unavailable.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") - @pytest.mark.skipif(check_device_drivers("gpu"), - reason="nvidia-smi not found") + reason="Language models currently failing for dynamic case" + ) + @pytest.mark.skipif( + check_device_drivers("gpu"), reason="nvidia-smi not found" + ) def test_module_dynamic_gpu(self): dynamic = True device = "gpu" @@ -63,8 +65,7 @@ class MiniLMTFModuleTest(unittest.TestCase): @pytest.mark.skip(reason="TF testing temporarily unavailable.") @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_static_vulkan(self): dynamic = False @@ -73,11 +74,11 @@ class MiniLMTFModuleTest(unittest.TestCase): @pytest.mark.skip(reason="TF testing temporarily unavailable.") @pytest.mark.xfail( - reason="Language models currently failing for dynamic case") + reason="Language models currently failing for dynamic case" + ) @pytest.mark.skipif( check_device_drivers("vulkan"), - reason= - "vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases" + reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases", ) def test_module_dynamic_vulkan(self): dynamic = True @@ -85,5 +86,5 @@ class MiniLMTFModuleTest(unittest.TestCase): self.module_tester.create_and_check_module(dynamic, device) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tank/tf/seq_classification.py b/tank/tf/seq_classification.py index bc117556..ccdb319f 100755 --- a/tank/tf/seq_classification.py +++ b/tank/tf/seq_classification.py @@ -6,12 +6,15 @@ from shark.parser import shark_args import argparse -seq_parser = argparse.ArgumentParser(description='Shark Sequence Classification.') +seq_parser = argparse.ArgumentParser( + description="Shark Sequence Classification." +) seq_parser.add_argument( "--hf_model_name", type=str, default="bert-base-uncased", - help="Hugging face model to run sequence classification.") + help="Hugging face model to run sequence classification.", +) seq_args, unknown = seq_parser.parse_known_args() @@ -25,45 +28,56 @@ inputs_signature = [ tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32), ] -# For supported models please see here: +# For supported models please see here: # https://huggingface.co/docs/transformers/model_doc/auto#transformers.TFAutoModelForSequenceClassification -def preprocess_input(text = "This is just used to compile the model"): + +def preprocess_input(text="This is just used to compile the model"): tokenizer = AutoTokenizer.from_pretrained(seq_args.hf_model_name) - inputs = tokenizer(text, - padding="max_length", - return_tensors="tf", - truncation=True, - max_length=MAX_SEQUENCE_LENGTH) + inputs = tokenizer( + text, + padding="max_length", + return_tensors="tf", + truncation=True, + max_length=MAX_SEQUENCE_LENGTH, + ) return inputs class SeqClassification(tf.Module): - def __init__(self, model_name): super(SeqClassification, self).__init__() self.m = TFAutoModelForSequenceClassification.from_pretrained( - model_name, output_attentions=False, num_labels=2) + model_name, output_attentions=False, num_labels=2 + ) self.m.predict = lambda x, y: self.m(input_ids=x, attention_mask=y)[0] @tf.function(input_signature=inputs_signature) def forward(self, input_ids, attention_mask): - return tf.math.softmax(self.m.predict(input_ids, attention_mask), - axis=-1) + return tf.math.softmax( + self.m.predict(input_ids, attention_mask), axis=-1 + ) if __name__ == "__main__": inputs = preprocess_input() shark_module = SharkInference( SeqClassification(seq_args.hf_model_name), - (inputs["input_ids"], inputs["attention_mask"])) + (inputs["input_ids"], inputs["attention_mask"]), + ) shark_module.set_frontend("tensorflow") shark_module.compile() print(f"Model has been successfully compiled on {shark_args.device}") while True: - input_text = input("Enter the text to classify (press q or nothing to exit): ") + input_text = input( + "Enter the text to classify (press q or nothing to exit): " + ) if not input_text or input_text == "q": break inputs = preprocess_input(input_text) - print(shark_module.forward((inputs["input_ids"], inputs["attention_mask"]))) + print( + shark_module.forward( + (inputs["input_ids"], inputs["attention_mask"]) + ) + ) diff --git a/tank/tflite/albert.py b/tank/tflite/albert.py index 619e3dcc..ef3dd476 100755 --- a/tank/tflite/albert.py +++ b/tank/tflite/albert.py @@ -13,26 +13,35 @@ def generate_inputs(input_details): args = [] args.append( - np.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + 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"])) + 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"])) + np.zeros( + shape=input_details[2]["shape"], dtype=input_details[2]["dtype"] + ) + ) return args -if __name__ == '__main__': - my_shark_importer = SharkImporter(model_path=model_path, - model_type="tflite", - model_source_hub="tfhub", - device="cpu", - dynamic=False, - jit_trace=True) +if __name__ == "__main__": + my_shark_importer = SharkImporter( + model_path=model_path, + model_type="tflite", + model_source_hub="tfhub", + device="cpu", + dynamic=False, + jit_trace=True, + ) # Case1: Use default inputs my_shark_importer.compile() shark_results = my_shark_importer.forward() diff --git a/tank/tflite/albert_lite_base/albert_lite_base_tflite_mlir_test.py b/tank/tflite/albert_lite_base/albert_lite_base_tflite_mlir_test.py index e44cbd2b..7f93133c 100644 --- a/tank/tflite/albert_lite_base/albert_lite_base_tflite_mlir_test.py +++ b/tank/tflite/albert_lite_base/albert_lite_base_tflite_mlir_test.py @@ -6,13 +6,12 @@ from shark.parser import shark_args class AlbertTfliteModuleTester: - def __init__( - self, - dynamic=False, - device="cpu", - save_mlir=False, - save_vmfb=False, + self, + dynamic=False, + device="cpu", + save_mlir=False, + save_vmfb=False, ): self.dynamic = dynamic self.device = device @@ -22,19 +21,23 @@ class AlbertTfliteModuleTester: def create_and_check_module(self): shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - self.shark_downloader = SharkDownloader(model_name="albert_lite_base", - tank_url="https://storage.googleapis.com/shark_tank", - local_tank_dir="./../gen_shark_tank/tflite", - model_type="tflite-tosa", - input_json="input.json", - input_type="int32") + self.shark_downloader = SharkDownloader( + model_name="albert_lite_base", + tank_url="https://storage.googleapis.com/shark_tank", + local_tank_dir="./../gen_shark_tank/tflite", + model_type="tflite-tosa", + input_json="input.json", + input_type="int32", + ) tflite_tosa_model = self.shark_downloader.get_mlir_file() inputs = self.shark_downloader.get_inputs() - self.shark_module = SharkInference(tflite_tosa_model, - inputs, - device=self.device, - dynamic=self.dynamic, - jit_trace=True) + self.shark_module = SharkInference( + tflite_tosa_model, + inputs, + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + ) self.shark_module.set_frontend("tflite-tosa") self.shark_module.compile() self.shark_module.forward(inputs) @@ -42,7 +45,6 @@ class AlbertTfliteModuleTester: class AlbertTfliteModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") @@ -58,7 +60,7 @@ class AlbertTfliteModuleTest(unittest.TestCase): self.module_tester.create_and_check_module() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() # module_tester = AlbertTfliteModuleTester() # module_tester.create_and_check_module() diff --git a/tank/tflite/albert_lite_base/albert_lite_base_tflite_test.py b/tank/tflite/albert_lite_base/albert_lite_base_tflite_test.py index e1e0c67f..4f463a9b 100644 --- a/tank/tflite/albert_lite_base/albert_lite_base_tflite_test.py +++ b/tank/tflite/albert_lite_base/albert_lite_base_tflite_test.py @@ -14,20 +14,27 @@ def generate_inputs(input_details): args = [] args.append( - np.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + 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"])) + 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"])) + np.zeros( + shape=input_details[2]["shape"], dtype=input_details[2]["dtype"] + ) + ) return args -class AlbertTfliteModuleTester: +class AlbertTfliteModuleTester: def __init__( self, dynamic=False, @@ -43,14 +50,16 @@ class AlbertTfliteModuleTester: def create_and_check_module(self): shark_args.save_mlir = self.save_mlir shark_args.save_vmfb = self.save_vmfb - my_shark_importer = SharkImporter(model_name="albert_lite_base", - # model_path=model_path, - model_type="tflite", - model_source_hub="tfhub", - device=self.device, - dynamic=self.dynamic, - jit_trace=True, - tank_url=None) + my_shark_importer = SharkImporter( + model_name="albert_lite_base", + # model_path=model_path, + model_type="tflite", + model_source_hub="tfhub", + device=self.device, + dynamic=self.dynamic, + jit_trace=True, + tank_url=None, + ) # Case1: Use default inputs my_shark_importer.compile() shark_results = my_shark_importer.forward() @@ -63,7 +72,6 @@ class AlbertTfliteModuleTester: class AlbertTfliteModuleTest(unittest.TestCase): - @pytest.fixture(autouse=True) def configure(self, pytestconfig): self.save_mlir = pytestconfig.getoption("save_mlir") @@ -78,7 +86,8 @@ class AlbertTfliteModuleTest(unittest.TestCase): self.module_tester.device = "cpu" self.module_tester.create_and_check_module() -if __name__ == '__main__': + +if __name__ == "__main__": # module_tester = AlbertTfliteModuleTester() # module_tester.create_and_check_module() unittest.main() diff --git a/tank/tflite/asr_conformer_test.py b/tank/tflite/asr_conformer_test.py index 5d2b10f4..88654b23 100644 --- a/tank/tflite/asr_conformer_test.py +++ b/tank/tflite/asr_conformer_test.py @@ -10,7 +10,6 @@ model_path = "https://tfhub.dev/neso613/lite-model/ASR_TFLite/pre_trained_models # Failure is due to dynamic shapes: # - Some improvements to tfl.strided_slice lowering are next steps class AsrConformerTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(AsrConformerTest, self).__init__(model_path, *args, **kwargs) @@ -18,5 +17,5 @@ class AsrConformerTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/bird_classifier_test.py b/tank/tflite/bird_classifier_test.py index ce099e2a..829bca8c 100644 --- a/tank/tflite/bird_classifier_test.py +++ b/tank/tflite/bird_classifier_test.py @@ -11,18 +11,21 @@ model_path = "https://tfhub.dev/google/lite-model/aiy/vision/classifier/birds_V1 class BirdClassifierTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(BirdClassifierTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(BirdClassifierTest, self).compare_results(iree_results, - tflite_results, details) + super(BirdClassifierTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) def generate_inputs(self, input_details): - img_path = "https://github.com/google-coral/test_data/raw/master/bird.bmp" + img_path = ( + "https://github.com/google-coral/test_data/raw/master/bird.bmp" + ) local_path = "/".join([self.workdir, "bird.bmp"]) urllib.request.urlretrieve(img_path, local_path) @@ -35,5 +38,5 @@ class BirdClassifierTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/cartoon_gan_test.py b/tank/tflite/cartoon_gan_test.py index 761b1941..ef102eb5 100644 --- a/tank/tflite/cartoon_gan_test.py +++ b/tank/tflite/cartoon_gan_test.py @@ -4,11 +4,12 @@ import absl.testing import test_util -model_path = "https://tfhub.dev/sayakpaul/lite-model/cartoongan/dr/1?lite-format=tflite" +model_path = ( + "https://tfhub.dev/sayakpaul/lite-model/cartoongan/dr/1?lite-format=tflite" +) class CartoonGanTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(CartoonGanTest, self).__init__(model_path, *args, **kwargs) @@ -16,5 +17,5 @@ class CartoonGanTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/craft_text_test.py b/tank/tflite/craft_text_test.py index c53b887d..f6f2bef8 100644 --- a/tank/tflite/craft_text_test.py +++ b/tank/tflite/craft_text_test.py @@ -10,17 +10,17 @@ model_path = "https://tfhub.dev/tulasiram58827/lite-model/craft-text-detector/dr # Failure: Resize lowering does not handle inferred dynamic shapes. Furthermore, the entire model # requires dynamic shape support. class CraftTextTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(CraftTextTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(CraftTextTest, self).compare_results(iree_results, tflite_results, - details) + super(CraftTextTest, self).compare_results( + iree_results, tflite_results, details + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/deeplab_v3_test.py b/tank/tflite/deeplab_v3_test.py index 331b285a..b01b17d8 100644 --- a/tank/tflite/deeplab_v3_test.py +++ b/tank/tflite/deeplab_v3_test.py @@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2?lit class DeepLabV3Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(DeepLabV3Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(DeepLabV3Test, self).compare_results(iree_results, tflite_results, - details) + super(DeepLabV3Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/densenet_test.py b/tank/tflite/densenet_test.py index 4e01d707..12563c4c 100644 --- a/tank/tflite/densenet_test.py +++ b/tank/tflite/densenet_test.py @@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/densenet/1/metadata/1?lite class DenseNetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(DenseNetTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(DenseNetTest, self).compare_results(iree_results, tflite_results, - details) + super(DenseNetTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-5).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-5).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/east_text_detector_test.py b/tank/tflite/east_text_detector_test.py index 3d444cef..9af041a7 100644 --- a/tank/tflite/east_text_detector_test.py +++ b/tank/tflite/east_text_detector_test.py @@ -8,28 +8,30 @@ model_path = "https://tfhub.dev/sayakpaul/lite-model/east-text-detector/dr/1?lit class EastTextDetectorTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(EastTextDetectorTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(EastTextDetectorTest, - self).compare_results(iree_results, tflite_results, details) + super(EastTextDetectorTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) # The second return is extremely noisy as it is not a binary classification. To handle we # check normalized correlation with an expectation of "close enough". iree_norm = numpy.sqrt(iree_results[1] * iree_results[1]) tflite_norm = numpy.sqrt(tflite_results[1] * tflite_results[1]) - correlation = numpy.average(iree_results[1] * tflite_results[1] / - iree_norm / tflite_norm) + correlation = numpy.average( + iree_results[1] * tflite_results[1] / iree_norm / tflite_norm + ) self.assertTrue(numpy.isclose(correlation, 1.0, atol=1e-2).all()) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/efficientnet_lite0_int8_test.py b/tank/tflite/efficientnet_lite0_int8_test.py index 72166c76..64e20df2 100644 --- a/tank/tflite/efficientnet_lite0_int8_test.py +++ b/tank/tflite/efficientnet_lite0_int8_test.py @@ -10,23 +10,25 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/efficientnet_l class EfficientnetLite0Int8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(EfficientnetLite0Int8Test, self).__init__(model_path, *args, - **kwargs) + super(EfficientnetLite0Int8Test, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(EfficientnetLite0Int8Test, - self).compare_results(iree_results, tflite_results, details) + super(EfficientnetLite0Int8Test, self).compare_results( + iree_results, tflite_results, details + ) # Dequantize outputs. - zero_point = details[0]['quantization_parameters']['zero_points'][0] - scale = details[0]['quantization_parameters']['scales'][0] + zero_point = details[0]["quantization_parameters"]["zero_points"][0] + scale = details[0]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results - zero_point) * scale dequantized_tflite_results = (tflite_results - zero_point) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=5e-3).all()) + numpy.isclose( + dequantized_iree_results, dequantized_tflite_results, atol=5e-3 + ).all() + ) def generate_inputs(self, input_details): return [imagenet_test_data.generate_input(self.workdir, input_details)] @@ -35,5 +37,5 @@ class EfficientnetLite0Int8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/efficientnet_lite0_test.py b/tank/tflite/efficientnet_lite0_test.py index 0c4cb0bc..eaf47fa6 100644 --- a/tank/tflite/efficientnet_lite0_test.py +++ b/tank/tflite/efficientnet_lite0_test.py @@ -10,25 +10,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/efficientnet_l class EfficientnetLite0Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(EfficientnetLite0Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(EfficientnetLite0Test, - self).compare_results(iree_results, tflite_results, details) + super(EfficientnetLite0Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/efficientnet_test.py b/tank/tflite/efficientnet_test.py index 62da5ee9..f97e4c6c 100644 --- a/tank/tflite/efficientnet_test.py +++ b/tank/tflite/efficientnet_test.py @@ -11,25 +11,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/efficientnet_2 class EfficientnetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(EfficientnetTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(EfficientnetTest, self).compare_results(iree_results, - tflite_results, details) + super(EfficientnetTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/gpt2_test.py b/tank/tflite/gpt2_test.py index 0a242bfd..20dedc6a 100644 --- a/tank/tflite/gpt2_test.py +++ b/tank/tflite/gpt2_test.py @@ -5,12 +5,13 @@ import absl.testing import numpy import test_util -model_path = "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-64.tflite" +model_path = ( + "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-64.tflite" +) # This test is a massive download and excluded due to causing timeouts. class GPT2Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(GPT2Test, self).__init__(model_path, *args, **kwargs) @@ -18,23 +19,29 @@ class GPT2Test(test_util.TFLiteModelTest): def generate_inputs(self, input_details): args = [] args.append( - numpy.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + numpy.random.randint( + low=0, + high=256, + size=input_details[0]["shape"], + dtype=input_details[0]["dtype"], + ) + ) return args def compare_results(self, iree_results, tflite_results, details): - super(GPT2Test, self).compare_results(iree_results, tflite_results, - details) + super(GPT2Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): self.assertTrue( - numpy.isclose(iree_results[i], tflite_results[i], - atol=5e-3).all()) + numpy.isclose( + iree_results[i], tflite_results[i], atol=5e-3 + ).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/image_stylization_test.py b/tank/tflite/image_stylization_test.py index 4086b7a7..68bc6375 100644 --- a/tank/tflite/image_stylization_test.py +++ b/tank/tflite/image_stylization_test.py @@ -9,23 +9,25 @@ model_path = "https://tfhub.dev/sayakpaul/lite-model/arbitrary-image-stylization # Failure is due to avg_pool2d. class ImageStylizationTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(ImageStylizationTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(ImageStylizationTest, - self).compare_results(iree_results, tflite_results, details) + super(ImageStylizationTest, self).compare_results( + iree_results, tflite_results, details + ) iree = iree_results[0].flatten().astype(numpy.single) tflite = tflite_results[0].flatten().astype(numpy.single) # Error is not tiny but appears close. self.assertTrue( - numpy.isclose(numpy.max(numpy.abs(iree - tflite)), 0.0, - atol=5e-2).all()) + numpy.isclose( + numpy.max(numpy.abs(iree - tflite)), 0.0, atol=5e-2 + ).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/inception_v4_test.py b/tank/tflite/inception_v4_test.py index fb6b4454..f08dd4b1 100644 --- a/tank/tflite/inception_v4_test.py +++ b/tank/tflite/inception_v4_test.py @@ -10,25 +10,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/inception_v4_2 class InceptionV4Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(InceptionV4Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(InceptionV4Test, self).compare_results(iree_results, - tflite_results, details) + super(InceptionV4Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/inception_v4_uint8_test.py b/tank/tflite/inception_v4_uint8_test.py index c52f1348..a5ed89fb 100644 --- a/tank/tflite/inception_v4_uint8_test.py +++ b/tank/tflite/inception_v4_uint8_test.py @@ -10,22 +10,23 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/inception_v4_2 class InceptionV4Uint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(InceptionV4Uint8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(InceptionV4Uint8Test, - self).compare_results(iree_results, tflite_results, details) + super(InceptionV4Uint8Test, self).compare_results( + iree_results, tflite_results, details + ) # Dequantize outputs. - zero_point = details[0]['quantization_parameters']['zero_points'][0] - scale = details[0]['quantization_parameters']['scales'][0] + zero_point = details[0]["quantization_parameters"]["zero_points"][0] + scale = details[0]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results - zero_point) * scale dequantized_tflite_results = (tflite_results - zero_point) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=5e-3).all()) + numpy.isclose( + dequantized_iree_results, dequantized_tflite_results, atol=5e-3 + ).all() + ) def generate_inputs(self, input_details): return [imagenet_test_data.generate_input(self.workdir, input_details)] @@ -34,5 +35,5 @@ class InceptionV4Uint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/keras_ocr_test.py b/tank/tflite/keras_ocr_test.py index 0dbbd99b..4ca359c6 100644 --- a/tank/tflite/keras_ocr_test.py +++ b/tank/tflite/keras_ocr_test.py @@ -8,7 +8,6 @@ model_path = "https://tfhub.dev/tulasiram58827/lite-model/keras-ocr/dr/2?lite-fo class KerasOCRTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(KerasOCRTest, self).__init__(model_path, *args, **kwargs) @@ -16,5 +15,5 @@ class KerasOCRTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/lightning_fp16_test.py b/tank/tflite/lightning_fp16_test.py index d42fbbb8..e4c514b8 100644 --- a/tank/tflite/lightning_fp16_test.py +++ b/tank/tflite/lightning_fp16_test.py @@ -10,7 +10,6 @@ model_path = "https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/t # Currently failing further in the linalg stack: # Bug related to linalg fusion. Collapsing dimension despite linalg index. class LightningFp16Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(LightningFp16Test, self).__init__(model_path, *args, **kwargs) @@ -18,5 +17,5 @@ class LightningFp16Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/lightning_i8_test.py b/tank/tflite/lightning_i8_test.py index 2f0b6b31..62c04418 100644 --- a/tank/tflite/lightning_i8_test.py +++ b/tank/tflite/lightning_i8_test.py @@ -13,7 +13,6 @@ model_path = "https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/t # Currently failing further in the linalg stack: # Invalid cast from ui8 to f32 TODO: make tfl.cast insert a rescale for ui8 class LightningI8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(LightningI8Test, self).__init__(model_path, *args, **kwargs) @@ -22,6 +21,7 @@ class LightningI8Test(test_util.TFLiteModelTest): # debug the source of numerical differences. def plot_results(self, iree_results, tflite_results, details): from matplotlib import pyplot as plt + local_path = "/".join([self.workdir, "person.jpg"]) im = numpy.array(Image.open(local_path)) @@ -38,24 +38,31 @@ class LightningI8Test(test_util.TFLiteModelTest): iree_y = iree_result[0, 0, :, 1] * height plt.imshow(im) - plt.scatter(tflite_y, tflite_x, label='tflite') - plt.scatter(iree_y, iree_x, label='iree') + plt.scatter(tflite_y, tflite_x, label="tflite") + plt.scatter(iree_y, iree_x, label="iree") plt.legend() plt.show() def compare_results(self, iree_results, tflite_results, details): - super(LightningI8Test, self).compare_results(iree_results, - tflite_results, details) + super(LightningI8Test, self).compare_results( + iree_results, tflite_results, details + ) # This value is a discretized location of the persons joints. If we are # *close* to the expected position we can consider this good enough. self.assertTrue( - numpy.isclose(iree_results[0][:, :, :, 0], - tflite_results[0][:, :, :, 0], - atol=25e-3).all()) + numpy.isclose( + iree_results[0][:, :, :, 0], + tflite_results[0][:, :, :, 0], + atol=25e-3, + ).all() + ) self.assertTrue( - numpy.isclose(iree_results[0][:, :, :, 1], - tflite_results[0][:, :, :, 1], - atol=25e-3).all()) + numpy.isclose( + iree_results[0][:, :, :, 1], + tflite_results[0][:, :, :, 1], + atol=25e-3, + ).all() + ) # self.plot_results(iree_results, tflite_results, details) def generate_inputs(self, input_details): @@ -72,5 +79,5 @@ class LightningI8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/lightning_test.py b/tank/tflite/lightning_test.py index 918be465..fd38182e 100644 --- a/tank/tflite/lightning_test.py +++ b/tank/tflite/lightning_test.py @@ -10,7 +10,6 @@ model_path = "https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/3 # Currently failing further in the linalg stack: # Fusion appears to produce an invalid IR. class LightningTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(LightningTest, self).__init__(model_path, *args, **kwargs) @@ -18,5 +17,5 @@ class LightningTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/lit.cfg.py b/tank/tflite/lit.cfg.py index a726b104..a7d7d7fc 100644 --- a/tank/tflite/lit.cfg.py +++ b/tank/tflite/lit.cfg.py @@ -10,37 +10,39 @@ import lit.llvm lit.llvm.initialize(lit_config, config) # name: The name of this test suite. -config.name = 'TFLITEHUB' +config.name = "TFLITEHUB" config.test_format = lit.formats.ShTest() # suffixes: A list of file extensions to treat as test files. -config.suffixes = ['.py'] +config.suffixes = [".py"] # test_source_root: The root path where tests are located. config.test_source_root = os.path.dirname(__file__) -#config.use_default_substitutions() +# config.use_default_substitutions() config.excludes = [ - 'coco_test_data.py', - 'imagenet_test_data.py', - 'lit.cfg.py', - 'lit.site.cfg.py', - 'manual_test.py', - 'squad_test_data.py', - 'test_util.py', + "coco_test_data.py", + "imagenet_test_data.py", + "lit.cfg.py", + "lit.site.cfg.py", + "manual_test.py", + "squad_test_data.py", + "test_util.py", ] -config.substitutions.extend([ - ('%PYTHON', sys.executable), -]) +config.substitutions.extend( + [ + ("%PYTHON", sys.executable), + ] +) -config.environment['PYTHONPATH'] = ":".join(sys.path) +config.environment["PYTHONPATH"] = ":".join(sys.path) project_root = os.path.dirname(os.path.dirname(__file__)) # Enable features based on -D FEATURES=hugetest,vulkan # syntax. -features_param = lit_config.params.get('FEATURES') +features_param = lit_config.params.get("FEATURES") if features_param: - config.available_features.update(features_param.split(',')) + config.available_features.update(features_param.split(",")) diff --git a/tank/tflite/magenta_test.py b/tank/tflite/magenta_test.py index 39250236..5feb9553 100644 --- a/tank/tflite/magenta_test.py +++ b/tank/tflite/magenta_test.py @@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/google/lite-model/magenta/arbitrary-image-styliz class MagentaTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MagentaTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MagentaTest, self).compare_results(iree_results, tflite_results, - details) + super(MagentaTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=2e-1).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=2e-1).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/manual_test.py b/tank/tflite/manual_test.py index 915b157e..04d8a4c5 100644 --- a/tank/tflite/manual_test.py +++ b/tank/tflite/manual_test.py @@ -2,23 +2,24 @@ import absl.flags import absl.testing import test_util -absl.flags.DEFINE_string('model', None, 'model path to execute') +absl.flags.DEFINE_string("model", None, "model path to execute") class ManualTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(ManualTest, self).__init__(absl.flags.FLAGS.model, *args, - **kwargs) + super(ManualTest, self).__init__( + absl.flags.FLAGS.model, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(ManualTest, self).compare_results(iree_results, tflite_results, - details) + super(ManualTest, self).compare_results( + iree_results, tflite_results, details + ) def test_compile_tflite(self): if self.model_path is not None: self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/midas_test.py b/tank/tflite/midas_test.py index 630d7360..cbc8f301 100644 --- a/tank/tflite/midas_test.py +++ b/tank/tflite/midas_test.py @@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/intel/lite-model/midas/v2_1_small/1/lite/1?lite- class MidasTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MidasTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MidasTest, self).compare_results(iree_results, tflite_results, - details) + super(MidasTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mirnet_test.py b/tank/tflite/mirnet_test.py index d0762866..2ffa4ec2 100644 --- a/tank/tflite/mirnet_test.py +++ b/tank/tflite/mirnet_test.py @@ -10,19 +10,20 @@ model_path = "https://tfhub.dev/sayakpaul/lite-model/mirnet-fixed/dr/1?lite-form # Note this one takes forever right now. Great for performance work! class MirnetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MirnetTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MirnetTest, self).compare_results(iree_results, tflite_results, - details) + super(MirnetTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=5e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=5e-3).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mnasnet_test.py b/tank/tflite/mnasnet_test.py index c3657e8b..2ceef7f2 100644 --- a/tank/tflite/mnasnet_test.py +++ b/tank/tflite/mnasnet_test.py @@ -8,7 +8,6 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/mnasnet_1.0_224/1/metadata class MnasnetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MnasnetTest, self).__init__(model_path, *args, **kwargs) @@ -16,5 +15,5 @@ class MnasnetTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilebert_edgetpu_s_float_test.py b/tank/tflite/mobilebert_edgetpu_s_float_test.py index 9f2d310b..e4ece1a8 100644 --- a/tank/tflite/mobilebert_edgetpu_s_float_test.py +++ b/tank/tflite/mobilebert_edgetpu_s_float_test.py @@ -8,41 +8,51 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-edg class MobileBertTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobileBertTest, self).__init__(model_path, *args, **kwargs) # Inputs modified to be useful mobilebert inputs. def generate_inputs(self, input_details): for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) args = [] args.append( - numpy.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + numpy.random.randint( + low=0, + high=256, + size=input_details[0]["shape"], + dtype=input_details[0]["dtype"], + ) + ) args.append( - numpy.ones(shape=input_details[1]["shape"], - dtype=input_details[1]["dtype"])) + numpy.ones( + shape=input_details[1]["shape"], dtype=input_details[1]["dtype"] + ) + ) args.append( - numpy.zeros(shape=input_details[2]["shape"], - dtype=input_details[2]["dtype"])) + numpy.zeros( + shape=input_details[2]["shape"], dtype=input_details[2]["dtype"] + ) + ) return args def compare_results(self, iree_results, tflite_results, details): - super(MobileBertTest, self).compare_results(iree_results, - tflite_results, details) + super(MobileBertTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-4).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-4).all() + ) self.assertTrue( - numpy.isclose(iree_results[1], tflite_results[1], atol=1e-4).all()) + numpy.isclose(iree_results[1], tflite_results[1], atol=1e-4).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilebert_edgetpu_s_quant_test.py b/tank/tflite/mobilebert_edgetpu_s_quant_test.py index 4d31b262..15073dfd 100644 --- a/tank/tflite/mobilebert_edgetpu_s_quant_test.py +++ b/tank/tflite/mobilebert_edgetpu_s_quant_test.py @@ -9,41 +9,51 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-edg class MobileBertTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobileBertTest, self).__init__(model_path, *args, **kwargs) # Inputs modified to be useful mobilebert inputs. def generate_inputs(self, input_details): for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) args = [] args.append( - numpy.random.randint(low=0, - high=256, - size=input_details[0]["shape"], - dtype=input_details[0]["dtype"])) + numpy.random.randint( + low=0, + high=256, + size=input_details[0]["shape"], + dtype=input_details[0]["dtype"], + ) + ) args.append( - numpy.ones(shape=input_details[1]["shape"], - dtype=input_details[1]["dtype"])) + numpy.ones( + shape=input_details[1]["shape"], dtype=input_details[1]["dtype"] + ) + ) args.append( - numpy.zeros(shape=input_details[2]["shape"], - dtype=input_details[2]["dtype"])) + numpy.zeros( + shape=input_details[2]["shape"], dtype=input_details[2]["dtype"] + ) + ) return args def compare_results(self, iree_results, tflite_results, details): - super(MobileBertTest, self).compare_results(iree_results, - tflite_results, details) + super(MobileBertTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all() + ) self.assertTrue( - numpy.isclose(iree_results[1], tflite_results[1], atol=1.0).all()) + numpy.isclose(iree_results[1], tflite_results[1], atol=1.0).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilebert_test.py b/tank/tflite/mobilebert_test.py index 504bb372..97b4ac7d 100644 --- a/tank/tflite/mobilebert_test.py +++ b/tank/tflite/mobilebert_test.py @@ -9,39 +9,45 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/mobilebert/1/metadata/1?li class MobileBertTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobileBertTest, self).__init__(model_path, *args, **kwargs) # Inputs modified to be useful mobilebert inputs. def generate_inputs(self, input_details): for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) - input_0 = np.asarray(squad_test_data._INPUT_WORD_ID, - dtype=input_details[0]["dtype"]) - input_1 = np.asarray(squad_test_data._INPUT_TYPE_ID, - dtype=input_details[1]["dtype"]) - input_2 = np.asarray(squad_test_data._INPUT_MASK, - dtype=input_details[2]["dtype"]) + input_0 = np.asarray( + squad_test_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"] + ) + input_1 = np.asarray( + squad_test_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"] + ) + input_2 = np.asarray( + squad_test_data._INPUT_MASK, dtype=input_details[2]["dtype"] + ) return [ input_0.reshape(input_details[0]["shape"]), input_1.reshape(input_details[1]["shape"]), - input_2.reshape(input_details[2]["shape"]) + input_2.reshape(input_details[2]["shape"]), ] def compare_results(self, iree_results, tflite_results, details): - super(MobileBertTest, self).compare_results(iree_results, - tflite_results, details) + super(MobileBertTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - np.isclose(iree_results[0], tflite_results[0], atol=1e-4).all()) + np.isclose(iree_results[0], tflite_results[0], atol=1e-4).all() + ) self.assertTrue( - np.isclose(iree_results[1], tflite_results[1], atol=1e-4).all()) + np.isclose(iree_results[1], tflite_results[1], atol=1e-4).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilebert_tf2_float_test.py b/tank/tflite/mobilebert_tf2_float_test.py index 07860a32..0a1b5032 100644 --- a/tank/tflite/mobilebert_tf2_float_test.py +++ b/tank/tflite/mobilebert_tf2_float_test.py @@ -9,39 +9,45 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-bas class MobileBertTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobileBertTest, self).__init__(model_path, *args, **kwargs) # Inputs modified to be useful mobilebert inputs. def generate_inputs(self, input_details): for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) - input_0 = np.asarray(squad_test_data._INPUT_WORD_ID, - dtype=input_details[0]["dtype"]) - input_1 = np.asarray(squad_test_data._INPUT_TYPE_ID, - dtype=input_details[1]["dtype"]) - input_2 = np.asarray(squad_test_data._INPUT_MASK, - dtype=input_details[2]["dtype"]) + input_0 = np.asarray( + squad_test_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"] + ) + input_1 = np.asarray( + squad_test_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"] + ) + input_2 = np.asarray( + squad_test_data._INPUT_MASK, dtype=input_details[2]["dtype"] + ) return [ input_0.reshape(input_details[0]["shape"]), input_1.reshape(input_details[1]["shape"]), - input_2.reshape(input_details[2]["shape"]) + input_2.reshape(input_details[2]["shape"]), ] def compare_results(self, iree_results, tflite_results, details): - super(MobileBertTest, self).compare_results(iree_results, - tflite_results, details) + super(MobileBertTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - np.isclose(iree_results[0], tflite_results[0], atol=1e-4).all()) + np.isclose(iree_results[0], tflite_results[0], atol=1e-4).all() + ) self.assertTrue( - np.isclose(iree_results[1], tflite_results[1], atol=1e-4).all()) + np.isclose(iree_results[1], tflite_results[1], atol=1e-4).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilebert_tf2_quant_test.py b/tank/tflite/mobilebert_tf2_quant_test.py index 746eda99..e3e038aa 100644 --- a/tank/tflite/mobilebert_tf2_quant_test.py +++ b/tank/tflite/mobilebert_tf2_quant_test.py @@ -9,40 +9,46 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-bas class MobileBertTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobileBertTest, self).__init__(model_path, *args, **kwargs) # Inputs modified to be useful mobilebert inputs. def generate_inputs(self, input_details): for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) - input_0 = np.asarray(squad_test_data._INPUT_WORD_ID, - dtype=input_details[0]["dtype"]) - input_1 = np.asarray(squad_test_data._INPUT_TYPE_ID, - dtype=input_details[1]["dtype"]) - input_2 = np.asarray(squad_test_data._INPUT_MASK, - dtype=input_details[2]["dtype"]) + input_0 = np.asarray( + squad_test_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"] + ) + input_1 = np.asarray( + squad_test_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"] + ) + input_2 = np.asarray( + squad_test_data._INPUT_MASK, dtype=input_details[2]["dtype"] + ) return [ input_0.reshape(input_details[0]["shape"]), input_1.reshape(input_details[1]["shape"]), - input_2.reshape(input_details[2]["shape"]) + input_2.reshape(input_details[2]["shape"]), ] def compare_results(self, iree_results, tflite_results, details): - super(MobileBertTest, self).compare_results(iree_results, - tflite_results, details) + super(MobileBertTest, self).compare_results( + iree_results, tflite_results, details + ) # We have confirmed in large scale accuracy tests that differences this large is acceptable. self.assertTrue( - np.isclose(iree_results[0], tflite_results[0], atol=5.0).all()) + np.isclose(iree_results[0], tflite_results[0], atol=5.0).all() + ) self.assertTrue( - np.isclose(iree_results[1], tflite_results[1], atol=5.0).all()) + np.isclose(iree_results[1], tflite_results[1], atol=5.0).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_ssd_quant_test.py b/tank/tflite/mobilenet_ssd_quant_test.py index e35dd5e2..d1b75b36 100644 --- a/tank/tflite/mobilenet_ssd_quant_test.py +++ b/tank/tflite/mobilenet_ssd_quant_test.py @@ -13,15 +13,16 @@ model_path = "https://storage.googleapis.com/iree-shared-files/models/ssd_mobile class MobilenetSsdQuantTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetSsdQuantTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetSsdQuantTest, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetSsdQuantTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all() + ) def generate_inputs(self, input_details): img_path = "https://github.com/google-coral/test_data/raw/master/grace_hopper.bmp" @@ -37,5 +38,5 @@ class MobilenetSsdQuantTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v1_test.py b/tank/tflite/mobilenet_v1_test.py index 60092b3d..78cf4973 100644 --- a/tank/tflite/mobilenet_v1_test.py +++ b/tank/tflite/mobilenet_v1_test.py @@ -9,25 +9,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v1_2 class MobilenetV1Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV1Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV1Test, self).compare_results(iree_results, - tflite_results, details) + super(MobilenetV1Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v1_uint8_test.py b/tank/tflite/mobilenet_v1_uint8_test.py index 430e72dd..faa77f2e 100644 --- a/tank/tflite/mobilenet_v1_uint8_test.py +++ b/tank/tflite/mobilenet_v1_uint8_test.py @@ -9,22 +9,23 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v1_2 class MobilenetV1Uint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV1Uint8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV1Uint8Test, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV1Uint8Test, self).compare_results( + iree_results, tflite_results, details + ) # Dequantize outputs. - zero_point = details[0]['quantization_parameters']['zero_points'][0] - scale = details[0]['quantization_parameters']['scales'][0] + zero_point = details[0]["quantization_parameters"]["zero_points"][0] + scale = details[0]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results - zero_point) * scale dequantized_tflite_results = (tflite_results - zero_point) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=5e-3).all()) + numpy.isclose( + dequantized_iree_results, dequantized_tflite_results, atol=5e-3 + ).all() + ) def generate_inputs(self, input_details): return [imagenet_test_data.generate_input(self.workdir, input_details)] @@ -33,5 +34,5 @@ class MobilenetV1Uint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v2_int8_test.py b/tank/tflite/mobilenet_v2_int8_test.py index 3ce3faca..601a5d2c 100644 --- a/tank/tflite/mobilenet_v2_int8_test.py +++ b/tank/tflite/mobilenet_v2_int8_test.py @@ -9,17 +9,18 @@ model_path = "https://storage.googleapis.com/tf_model_garden/vision/mobilenet/v2 class MobilenetV2Int8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV2Int8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV2Int8Test, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV2Int8Test, self).compare_results( + iree_results, tflite_results, details + ) # Although this a quantized model, inputs and outputs are in float. # The difference here is quite high for a dequantized output. self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=0.5).all()) + numpy.isclose(iree_results, tflite_results, atol=0.5).all() + ) # Make sure the predicted class is the same. iree_predicted_class = numpy.argmax(iree_results[0][0]) @@ -29,12 +30,12 @@ class MobilenetV2Int8Test(test_util.TFLiteModelTest): def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v2_test.py b/tank/tflite/mobilenet_v2_test.py index 6b219b18..f37436eb 100644 --- a/tank/tflite/mobilenet_v2_test.py +++ b/tank/tflite/mobilenet_v2_test.py @@ -9,25 +9,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_1 class MobilenetV2Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV2Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV2Test, self).compare_results(iree_results, - tflite_results, details) + super(MobilenetV2Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v2_uint8_test.py b/tank/tflite/mobilenet_v2_uint8_test.py index 4126e1bd..40f3189b 100644 --- a/tank/tflite/mobilenet_v2_uint8_test.py +++ b/tank/tflite/mobilenet_v2_uint8_test.py @@ -9,22 +9,23 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_2 class MobilenetV2Uint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV2Uint8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV2Uint8Test, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV2Uint8Test, self).compare_results( + iree_results, tflite_results, details + ) # Dequantize outputs. - zero_point = details[0]['quantization_parameters']['zero_points'][0] - scale = details[0]['quantization_parameters']['scales'][0] + zero_point = details[0]["quantization_parameters"]["zero_points"][0] + scale = details[0]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results - zero_point) * scale dequantized_tflite_results = (tflite_results - zero_point) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=5e-3).all()) + numpy.isclose( + dequantized_iree_results, dequantized_tflite_results, atol=5e-3 + ).all() + ) def generate_inputs(self, input_details): return [imagenet_test_data.generate_input(self.workdir, input_details)] @@ -33,5 +34,5 @@ class MobilenetV2Uint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v3-large_test.py b/tank/tflite/mobilenet_v3-large_test.py index 42ad7e4e..ab8e38b9 100644 --- a/tank/tflite/mobilenet_v3-large_test.py +++ b/tank/tflite/mobilenet_v3-large_test.py @@ -9,25 +9,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v3-l class MobilenetV3LargeTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV3LargeTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV3LargeTest, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV3LargeTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=1e-4).all()) + numpy.isclose(iree_results, tflite_results, atol=1e-4).all() + ) def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v3-large_uint8_test.py b/tank/tflite/mobilenet_v3-large_uint8_test.py index 8424e301..8ed48385 100644 --- a/tank/tflite/mobilenet_v3-large_uint8_test.py +++ b/tank/tflite/mobilenet_v3-large_uint8_test.py @@ -9,23 +9,25 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v3-l class MobilenetV3LargeUint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(MobilenetV3LargeUint8Test, self).__init__(model_path, *args, - **kwargs) + super(MobilenetV3LargeUint8Test, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV3LargeUint8Test, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV3LargeUint8Test, self).compare_results( + iree_results, tflite_results, details + ) # Dequantize outputs. - zero_point = details[0]['quantization_parameters']['zero_points'][0] - scale = details[0]['quantization_parameters']['scales'][0] + zero_point = details[0]["quantization_parameters"]["zero_points"][0] + scale = details[0]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results - zero_point) * scale dequantized_tflite_results = (tflite_results - zero_point) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=5e-3).all()) + numpy.isclose( + dequantized_iree_results, dequantized_tflite_results, atol=5e-3 + ).all() + ) def generate_inputs(self, input_details): return [imagenet_test_data.generate_input(self.workdir, input_details)] @@ -34,5 +36,5 @@ class MobilenetV3LargeUint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/mobilenet_v35_int8_test.py b/tank/tflite/mobilenet_v35_int8_test.py index d53457a5..ed4dd3d3 100644 --- a/tank/tflite/mobilenet_v35_int8_test.py +++ b/tank/tflite/mobilenet_v35_int8_test.py @@ -9,16 +9,17 @@ model_path = "https://storage.googleapis.com/tf_model_garden/vision/mobilenet/v3 class MobilenetV35Int8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MobilenetV35Int8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(MobilenetV35Int8Test, - self).compare_results(iree_results, tflite_results, details) + super(MobilenetV35Int8Test, self).compare_results( + iree_results, tflite_results, details + ) # The difference here is quite high for a dequantized output. self.assertTrue( - numpy.isclose(iree_results, tflite_results, atol=0.5).all()) + numpy.isclose(iree_results, tflite_results, atol=0.5).all() + ) # Make sure the predicted class is the same. iree_predicted_class = numpy.argmax(iree_results[0][0]) @@ -28,12 +29,12 @@ class MobilenetV35Int8Test(test_util.TFLiteModelTest): def generate_inputs(self, input_details): inputs = imagenet_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/nasnet_test.py b/tank/tflite/nasnet_test.py index 88e296be..bfa02782 100644 --- a/tank/tflite/nasnet_test.py +++ b/tank/tflite/nasnet_test.py @@ -9,7 +9,6 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/nasnet/large/1/default/1?l class MnasnetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(MnasnetTest, self).__init__(model_path, *args, **kwargs) @@ -17,5 +16,5 @@ class MnasnetTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/person_detect_test.py b/tank/tflite/person_detect_test.py index e56d8dca..06b0c041 100644 --- a/tank/tflite/person_detect_test.py +++ b/tank/tflite/person_detect_test.py @@ -11,27 +11,32 @@ model_path = "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea2 class PersonDetectTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(PersonDetectTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(PersonDetectTest, self).compare_results(iree_results, - tflite_results, details) + super(PersonDetectTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) # TFLite is broken with this model so we hardcode the input/output details. def setup_tflite(self): - self.input_details = [{ - "shape": [1, 96, 96, 1], - "dtype": numpy.int8, - "index": 0, - }] - self.output_details = [{ - "shape": [1, 2], - "dtype": numpy.int8, - }] + self.input_details = [ + { + "shape": [1, 96, 96, 1], + "dtype": numpy.int8, + "index": 0, + } + ] + self.output_details = [ + { + "shape": [1, 2], + "dtype": numpy.int8, + } + ] # The input has known expected values. We hardcode this value. def invoke_tflite(self, args): @@ -43,8 +48,9 @@ class PersonDetectTest(test_util.TFLiteModelTest): urllib.request.urlretrieve(img_path, local_path) shape = input_details[0]["shape"] - im = numpy.array(Image.open(local_path).resize( - (shape[1], shape[2]))).astype(input_details[0]["dtype"]) + im = numpy.array( + Image.open(local_path).resize((shape[1], shape[2])) + ).astype(input_details[0]["dtype"]) args = [im.reshape(shape)] return args @@ -52,5 +58,5 @@ class PersonDetectTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/posenet_test.py b/tank/tflite/posenet_test.py index 7183e17b..5dd5579b 100644 --- a/tank/tflite/posenet_test.py +++ b/tank/tflite/posenet_test.py @@ -8,25 +8,29 @@ model_path = "https://storage.googleapis.com/download.tensorflow.org/models/tfli class PoseTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(PoseTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(PoseTest, self).compare_results(iree_results, tflite_results, - details) + super(PoseTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all() + ) self.assertTrue( - numpy.isclose(iree_results[1], tflite_results[1], atol=1e-2).all()) + numpy.isclose(iree_results[1], tflite_results[1], atol=1e-2).all() + ) self.assertTrue( - numpy.isclose(iree_results[2], tflite_results[2], atol=1e-2).all()) + numpy.isclose(iree_results[2], tflite_results[2], atol=1e-2).all() + ) self.assertTrue( - numpy.isclose(iree_results[3], tflite_results[3], atol=1e-3).all()) + numpy.isclose(iree_results[3], tflite_results[3], atol=1e-3).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/resnet_50_int8_test.py b/tank/tflite/resnet_50_int8_test.py index 183c844a..7d7b8af1 100644 --- a/tank/tflite/resnet_50_int8_test.py +++ b/tank/tflite/resnet_50_int8_test.py @@ -8,19 +8,20 @@ model_path = "https://storage.googleapis.com/tf_model_garden/vision/resnet50_ima class ResNet50Int8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(ResNet50Int8Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(ResNet50Int8Test, self).compare_results(iree_results, - tflite_results, details) + super(ResNet50Int8Test, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1.0).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/rosetta_test.py b/tank/tflite/rosetta_test.py index fc6a0535..c38da51e 100644 --- a/tank/tflite/rosetta_test.py +++ b/tank/tflite/rosetta_test.py @@ -10,19 +10,20 @@ model_path = "https://tfhub.dev/tulasiram58827/lite-model/rosetta/dr/1?lite-form # tfl.padv2 cannot be lowered to tosa.pad. May be possible to switch tosa.concat class RosettaTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(RosettaTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(RosettaTest, self).compare_results(iree_results, tflite_results, - details) + super(RosettaTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=5e-3).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=5e-3).all() + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/spice_test.py b/tank/tflite/spice_test.py index 51978c76..8b4fe852 100644 --- a/tank/tflite/spice_test.py +++ b/tank/tflite/spice_test.py @@ -12,7 +12,6 @@ model_path = "https://tfhub.dev/google/lite-model/spice/1?lite-format=tflite" # 1. Multiple unsupported dynamic operations (tfl.stride, range, gather). # 2. Static version blocked by tfl.range not having a lowering for static fixed shapes. class SpiceTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(SpiceTest, self).__init__(model_path, *args, **kwargs) @@ -20,5 +19,5 @@ class SpiceTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/squad_test_data.py b/tank/tflite/squad_test_data.py index 20a73f3a..d624d51f 100644 --- a/tank/tflite/squad_test_data.py +++ b/tank/tflite/squad_test_data.py @@ -1,58 +1,1159 @@ # An example input combination from the Squad 1.1 dataset. _INPUT_WORD_ID = [ - 101, 2129, 2116, 19576, 2015, 2106, 3854, 4679, 2486, 1029, 102, 1996, - 14169, 2165, 2019, 2220, 2599, 1999, 3565, 4605, 2753, 1998, 2196, 11145, - 1012, 8446, 2001, 3132, 2011, 7573, 1005, 1055, 3639, 1010, 2029, 14159, - 2032, 2698, 2335, 1998, 3140, 2032, 2046, 2093, 20991, 2015, 1010, 2164, - 1037, 19576, 2029, 2027, 6757, 2005, 1037, 7921, 1012, 7573, 15674, 3854, - 4679, 2001, 2315, 3565, 4605, 12041, 1010, 3405, 2274, 3948, 10455, 1010, - 1016, 13714, 14918, 1010, 1998, 2048, 3140, 19576, 2015, 1012, 102, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + 101, + 2129, + 2116, + 19576, + 2015, + 2106, + 3854, + 4679, + 2486, + 1029, + 102, + 1996, + 14169, + 2165, + 2019, + 2220, + 2599, + 1999, + 3565, + 4605, + 2753, + 1998, + 2196, + 11145, + 1012, + 8446, + 2001, + 3132, + 2011, + 7573, + 1005, + 1055, + 3639, + 1010, + 2029, + 14159, + 2032, + 2698, + 2335, + 1998, + 3140, + 2032, + 2046, + 2093, + 20991, + 2015, + 1010, + 2164, + 1037, + 19576, + 2029, + 2027, + 6757, + 2005, + 1037, + 7921, + 1012, + 7573, + 15674, + 3854, + 4679, + 2001, + 2315, + 3565, + 4605, + 12041, + 1010, + 3405, + 2274, + 3948, + 10455, + 1010, + 1016, + 13714, + 14918, + 1010, + 1998, + 2048, + 3140, + 19576, + 2015, + 1012, + 102, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, ] _INPUT_TYPE_ID = [ - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0 + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, ] _INPUT_MASK = [ - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0 + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, ] diff --git a/tank/tflite/squeezenet_test.py b/tank/tflite/squeezenet_test.py index 452e9727..b53113bc 100644 --- a/tank/tflite/squeezenet_test.py +++ b/tank/tflite/squeezenet_test.py @@ -7,17 +7,17 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/squeezenet/1/default/1?lit class SqueezeNetTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(SqueezeNetTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(SqueezeNetTest, self).compare_results(iree_results, - tflite_results, details) + super(SqueezeNetTest, self).compare_results( + iree_results, tflite_results, details + ) def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v1_test.py b/tank/tflite/ssd_mobilenet_v1_test.py index 910d983b..595d733b 100644 --- a/tank/tflite/ssd_mobilenet_v1_test.py +++ b/tank/tflite/ssd_mobilenet_v1_test.py @@ -9,27 +9,29 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV1Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(SsdMobilenetV1Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV1Test, self).compare_results(iree_results, - tflite_results, details) + super(SsdMobilenetV1Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): self.assertTrue( - numpy.isclose(iree_results[i], tflite_results[i], - atol=1e-4).all()) + numpy.isclose( + iree_results[i], tflite_results[i], atol=1e-4 + ).all() + ) def generate_inputs(self, input_details): inputs = coco_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v1_uint8_test.py b/tank/tflite/ssd_mobilenet_v1_uint8_test.py index 2a047270..70fc05b6 100644 --- a/tank/tflite/ssd_mobilenet_v1_uint8_test.py +++ b/tank/tflite/ssd_mobilenet_v1_uint8_test.py @@ -9,25 +9,30 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV1Uint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(SsdMobilenetV1Uint8Test, self).__init__(model_path, *args, - **kwargs) + super(SsdMobilenetV1Uint8Test, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV1Uint8Test, - self).compare_results(iree_results, tflite_results, details) + super(SsdMobilenetV1Uint8Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): # Dequantize outputs. - zero_point = details[i]['quantization_parameters']['zero_points'][0] - scale = details[i]['quantization_parameters']['scales'][0] + zero_point = details[i]["quantization_parameters"]["zero_points"][0] + scale = details[i]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results[i] - zero_point) * scale - dequantized_tflite_results = (tflite_results[i] - - zero_point) * scale + dequantized_tflite_results = ( + tflite_results[i] - zero_point + ) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=0.1).all()) + numpy.isclose( + dequantized_iree_results, + dequantized_tflite_results, + atol=0.1, + ).all() + ) def generate_inputs(self, input_details): return [coco_test_data.generate_input(self.workdir, input_details)] @@ -36,5 +41,5 @@ class SsdMobilenetV1Uint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v2_fpnlite_test.py b/tank/tflite/ssd_mobilenet_v2_fpnlite_test.py index 20d3a301..6ade24ae 100644 --- a/tank/tflite/ssd_mobilenet_v2_fpnlite_test.py +++ b/tank/tflite/ssd_mobilenet_v2_fpnlite_test.py @@ -10,28 +10,31 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV2FpnliteTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(SsdMobilenetV2FpnliteTest, self).__init__(model_path, *args, - **kwargs) + super(SsdMobilenetV2FpnliteTest, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV2FpnliteTest, - self).compare_results(iree_results, tflite_results, details) + super(SsdMobilenetV2FpnliteTest, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): self.assertTrue( - numpy.isclose(iree_results[i], tflite_results[i], - atol=1e-4).all()) + numpy.isclose( + iree_results[i], tflite_results[i], atol=1e-4 + ).all() + ) def generate_inputs(self, input_details): inputs = coco_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v2_fpnlite_uint8_test.py b/tank/tflite/ssd_mobilenet_v2_fpnlite_uint8_test.py index 88598f03..40f46642 100644 --- a/tank/tflite/ssd_mobilenet_v2_fpnlite_uint8_test.py +++ b/tank/tflite/ssd_mobilenet_v2_fpnlite_uint8_test.py @@ -10,25 +10,30 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV2FpnliteUint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(SsdMobilenetV2FpnliteUint8Test, - self).__init__(model_path, *args, **kwargs) + super(SsdMobilenetV2FpnliteUint8Test, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV2FpnliteUint8Test, - self).compare_results(iree_results, tflite_results, details) + super(SsdMobilenetV2FpnliteUint8Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): # Dequantize outputs. - zero_point = details[i]['quantization_parameters']['zero_points'][0] - scale = details[i]['quantization_parameters']['scales'][0] + zero_point = details[i]["quantization_parameters"]["zero_points"][0] + scale = details[i]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results[i] - zero_point) * scale - dequantized_tflite_results = (tflite_results[i] - - zero_point) * scale + dequantized_tflite_results = ( + tflite_results[i] - zero_point + ) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=0.1).all()) + numpy.isclose( + dequantized_iree_results, + dequantized_tflite_results, + atol=0.1, + ).all() + ) def generate_inputs(self, input_details): return [coco_test_data.generate_input(self.workdir, input_details)] @@ -37,5 +42,5 @@ class SsdMobilenetV2FpnliteUint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v2_int8_test.py b/tank/tflite/ssd_mobilenet_v2_int8_test.py index 3ef57d42..7b83ea09 100644 --- a/tank/tflite/ssd_mobilenet_v2_int8_test.py +++ b/tank/tflite/ssd_mobilenet_v2_int8_test.py @@ -10,24 +10,28 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV2Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(SsdMobilenetV2Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV2Test, self).compare_results(iree_results, - tflite_results, details) + super(SsdMobilenetV2Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): # Dequantize outputs. - zero_point = details[i]['quantization_parameters']['zero_points'][0] - scale = details[i]['quantization_parameters']['scales'][0] + zero_point = details[i]["quantization_parameters"]["zero_points"][0] + scale = details[i]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results[i] - zero_point) * scale - dequantized_tflite_results = (tflite_results[i] - - zero_point) * scale + dequantized_tflite_results = ( + tflite_results[i] - zero_point + ) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=0.1).all()) + numpy.isclose( + dequantized_iree_results, + dequantized_tflite_results, + atol=0.1, + ).all() + ) def generate_inputs(self, input_details): inputs = coco_test_data.generate_input(self.workdir, input_details) @@ -39,5 +43,5 @@ class SsdMobilenetV2Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_mobilenet_v2_test.py b/tank/tflite/ssd_mobilenet_v2_test.py index 0e12ceae..8573d005 100644 --- a/tank/tflite/ssd_mobilenet_v2_test.py +++ b/tank/tflite/ssd_mobilenet_v2_test.py @@ -10,27 +10,29 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_mobilenet_ class SsdMobilenetV2Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(SsdMobilenetV2Test, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(SsdMobilenetV2Test, self).compare_results(iree_results, - tflite_results, details) + super(SsdMobilenetV2Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): self.assertTrue( - numpy.isclose(iree_results[i], tflite_results[i], - atol=1e-4).all()) + numpy.isclose( + iree_results[i], tflite_results[i], atol=1e-4 + ).all() + ) def generate_inputs(self, input_details): inputs = coco_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_spaghettinet_large_test.py b/tank/tflite/ssd_spaghettinet_large_test.py index afeb5fe3..8a6367ec 100644 --- a/tank/tflite/ssd_spaghettinet_large_test.py +++ b/tank/tflite/ssd_spaghettinet_large_test.py @@ -10,30 +10,33 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_spaghettin class SsdSpaghettinetLargeTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(SsdSpaghettinetLargeTest, self).__init__(model_path, *args, - **kwargs) + super(SsdSpaghettinetLargeTest, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(SsdSpaghettinetLargeTest, - self).compare_results(iree_results, tflite_results, details) + super(SsdSpaghettinetLargeTest, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): - print('iree_results: ' + str(iree_results[i])) - print('tflite_results: ' + str(tflite_results[i])) + print("iree_results: " + str(iree_results[i])) + print("tflite_results: " + str(tflite_results[i])) self.assertTrue( - numpy.isclose(iree_results[i], tflite_results[i], - atol=1e-4).all()) + numpy.isclose( + iree_results[i], tflite_results[i], atol=1e-4 + ).all() + ) def generate_inputs(self, input_details): inputs = coco_test_data.generate_input(self.workdir, input_details) # Normalize inputs to [-1, 1]. - inputs = (inputs.astype('float32') / 127.5) - 1 + inputs = (inputs.astype("float32") / 127.5) - 1 return [inputs] def test_compile_tflite(self): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/ssd_spaghettinet_large_uint8_test.py b/tank/tflite/ssd_spaghettinet_large_uint8_test.py index 85554ebb..b6247596 100644 --- a/tank/tflite/ssd_spaghettinet_large_uint8_test.py +++ b/tank/tflite/ssd_spaghettinet_large_uint8_test.py @@ -10,25 +10,30 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/ssd_spaghettin class SsdSpaghettinetLargeUint8Test(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): - super(SsdSpaghettinetLargeUint8Test, - self).__init__(model_path, *args, **kwargs) + super(SsdSpaghettinetLargeUint8Test, self).__init__( + model_path, *args, **kwargs + ) def compare_results(self, iree_results, tflite_results, details): - super(SsdSpaghettinetLargeUint8Test, - self).compare_results(iree_results, tflite_results, details) + super(SsdSpaghettinetLargeUint8Test, self).compare_results( + iree_results, tflite_results, details + ) for i in range(len(iree_results)): # Dequantize outputs. - zero_point = details[i]['quantization_parameters']['zero_points'][0] - scale = details[i]['quantization_parameters']['scales'][0] + zero_point = details[i]["quantization_parameters"]["zero_points"][0] + scale = details[i]["quantization_parameters"]["scales"][0] dequantized_iree_results = (iree_results[i] - zero_point) * scale - dequantized_tflite_results = (tflite_results[i] - - zero_point) * scale + dequantized_tflite_results = ( + tflite_results[i] - zero_point + ) * scale self.assertTrue( - numpy.isclose(dequantized_iree_results, - dequantized_tflite_results, - atol=0.1).all()) + numpy.isclose( + dequantized_iree_results, + dequantized_tflite_results, + atol=0.1, + ).all() + ) def generate_inputs(self, input_details): return [coco_test_data.generate_input(self.workdir, input_details)] @@ -37,5 +42,5 @@ class SsdSpaghettinetLargeUint8Test(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main() diff --git a/tank/tflite/test_util.py b/tank/tflite/test_util.py index a9eee2c2..953c868a 100644 --- a/tank/tflite/test_util.py +++ b/tank/tflite/test_util.py @@ -19,20 +19,19 @@ import time import urllib.request targets = { - 'dylib': 'dylib-llvm-aot', - 'vulkan': 'vulkan-spirv', + "dylib": "dylib-llvm-aot", + "vulkan": "vulkan-spirv", } configs = { - 'dylib': 'dylib', - 'vulkan': 'vulkan', + "dylib": "dylib", + "vulkan": "vulkan", } -absl.flags.DEFINE_string('config', 'dylib', 'model path to execute') +absl.flags.DEFINE_string("config", "dylib", "model path to execute") class TFLiteModelTest(testing.absltest.TestCase): - def __init__(self, model_path, *args, **kwargs): super(TFLiteModelTest, self).__init__(*args, **kwargs) self.model_path = model_path @@ -41,30 +40,35 @@ class TFLiteModelTest(testing.absltest.TestCase): if self.model_path is None: return exe_basename = os.path.basename(sys.argv[0]) - self.workdir = os.path.join(os.path.dirname(__file__), "tmp", - exe_basename) + self.workdir = os.path.join( + os.path.dirname(__file__), "tmp", exe_basename + ) print(f"TMP_DIR = {self.workdir}") os.makedirs(self.workdir, exist_ok=True) - self.tflite_file = '/'.join([self.workdir, 'model.tflite']) - self.tflite_ir = '/'.join([self.workdir, 'tflite.mlir']) - self.iree_ir = '/'.join([self.workdir, 'tosa.mlir']) + self.tflite_file = "/".join([self.workdir, "model.tflite"]) + self.tflite_ir = "/".join([self.workdir, "tflite.mlir"]) + self.iree_ir = "/".join([self.workdir, "tosa.mlir"]) if os.path.exists(self.model_path): self.tflite_file = self.model_path else: urllib.request.urlretrieve(self.model_path, self.tflite_file) - self.binary = '/'.join([self.workdir, 'module.bytecode']) + self.binary = "/".join([self.workdir, "module.bytecode"]) def generate_inputs(self, input_details): args = [] for input in input_details: - absl.logging.info("\t%s, %s", str(input["shape"]), - input["dtype"].__name__) + absl.logging.info( + "\t%s, %s", str(input["shape"]), input["dtype"].__name__ + ) args.append(np.zeros(shape=input["shape"], dtype=input["dtype"])) return args def compare_results(self, iree_results, tflite_results, details): - self.assertEqual(len(iree_results), len(tflite_results), - "Number of results do not match") + self.assertEqual( + len(iree_results), + len(tflite_results), + "Number of results do not match", + ) for i in range(len(details)): iree_result = iree_results[i] @@ -78,14 +82,15 @@ class TFLiteModelTest(testing.absltest.TestCase): def setup_tflite(self): absl.logging.info("Setting up tflite interpreter") self.tflite_interpreter = tf.lite.Interpreter( - model_path=self.tflite_file) + model_path=self.tflite_file + ) self.tflite_interpreter.allocate_tensors() self.input_details = self.tflite_interpreter.get_input_details() self.output_details = self.tflite_interpreter.get_output_details() def setup_iree(self): absl.logging.info("Setting up iree runtime") - with open(self.binary, 'rb') as f: + with open(self.binary, "rb") as f: config = iree_rt.Config(configs[absl.flags.FLAGS.config]) self.iree_context = iree_rt.SystemContext(config=config) vm_module = iree_rt.VmModule.from_flatbuffer(f.read()) @@ -93,8 +98,9 @@ class TFLiteModelTest(testing.absltest.TestCase): def invoke_tflite(self, args): for i, input in enumerate(args): - self.tflite_interpreter.set_tensor(self.input_details[i]['index'], - input) + self.tflite_interpreter.set_tensor( + self.input_details[i]["index"], input + ) start = time.perf_counter() self.tflite_interpreter.invoke() end = time.perf_counter() @@ -103,7 +109,9 @@ class TFLiteModelTest(testing.absltest.TestCase): for output_detail in self.output_details: tflite_results.append( np.array( - self.tflite_interpreter.get_tensor(output_detail['index']))) + self.tflite_interpreter.get_tensor(output_detail["index"]) + ) + ) for i in range(len(self.output_details)): dtype = self.output_details[i]["dtype"] @@ -131,7 +139,8 @@ class TFLiteModelTest(testing.absltest.TestCase): save_temp_tfl_input=self.tflite_ir, save_temp_iree_input=self.iree_ir, target_backends=[targets[absl.flags.FLAGS.config]], - import_only=False) + import_only=False, + ) self.setup_tflite() self.setup_iree() diff --git a/tank/tflite/visual_wake_words_i8_test.py b/tank/tflite/visual_wake_words_i8_test.py index dce92ac3..bdb6f639 100644 --- a/tank/tflite/visual_wake_words_i8_test.py +++ b/tank/tflite/visual_wake_words_i8_test.py @@ -14,15 +14,16 @@ model_path = "https://github.com/mlcommons/tiny/raw/0b04bcd402ee28f84e79fa86d8bb # and there is not currently supported. Flatbuffer was modified to use static # shapes and was otherwise numerically correct. class VisualWakeWordsTest(test_util.TFLiteModelTest): - def __init__(self, *args, **kwargs): super(VisualWakeWordsTest, self).__init__(model_path, *args, **kwargs) def compare_results(self, iree_results, tflite_results, details): - super(VisualWakeWordsTest, - self).compare_results(iree_results, tflite_results, details) + super(VisualWakeWordsTest, self).compare_results( + iree_results, tflite_results, details + ) self.assertTrue( - numpy.isclose(iree_results[0], tflite_results[0], atol=1).all()) + numpy.isclose(iree_results[0], tflite_results[0], atol=1).all() + ) def generate_inputs(self, input_details): img_path = "https://github.com/tensorflow/tflite-micro/raw/main/tensorflow/lite/micro/examples/person_detection/testdata/person.bmp" @@ -32,8 +33,10 @@ class VisualWakeWordsTest(test_util.TFLiteModelTest): shape = input_details[0]["shape"] input_type = input_details[0]["dtype"] im = numpy.array( - Image.open(local_path).resize( - (shape[1], shape[2])).convert(mode="RGB")) + Image.open(local_path) + .resize((shape[1], shape[2])) + .convert(mode="RGB") + ) args = [im.reshape(shape).astype(input_type)] return args @@ -41,5 +44,5 @@ class VisualWakeWordsTest(test_util.TFLiteModelTest): self.compile_and_execute() -if __name__ == '__main__': +if __name__ == "__main__": absl.testing.absltest.main()