github/AMD-SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/AMD-SHARK-Studio.git synced 2026-04-03 03:00:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

Add shark_importer for torch_models. (#183)

Browse Source 
All the torch_models are imported to gs::shark_tank.
Scripts have been updated.
This commit is contained in:
Prashant Kumar
2022-07-13 09:08:19 +05:30
committed by GitHub
parent 2e22d0b690
commit 0dcf387089
167 changed files with 2361 additions and 864 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/test-models.yml vendored
View File

@@ -65,7 +65,7 @@ jobs:
run: |
# black format check
black --version
black --line-length 127 --check .
black --line-length 79 --check .
# stop the build if there are Python syntax errors or undefined names
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

4
benchmarks/hf_model_benchmark.py
View File

@@ -13,7 +13,9 @@ 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)

16
benchmarks/hf_transformer.py
View File

@@ -56,7 +56,9 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
):
self.device = device if device is not None else shark_args.device
if self.device == "gpu":
raise ValueError("Currently GPU Benchmarking is not supported due to OOM from ORT.")
raise ValueError(
"Currently GPU Benchmarking is not supported due to OOM from ORT."
)
self.model_name = model_name
model = HuggingFaceLanguage(model_name)
SharkBenchmarkRunner.__init__(
@@ -93,7 +95,9 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
cache_dir,
verbose,
)
print(f"ONNX Pytorch-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}")
print(
f"ONNX Pytorch-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)
# TODO: Currently non-functional due to TF runtime error. There might be some issue with, initializing TF.
def benchmark_tf(self, inputs):
@@ -118,7 +122,9 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
cache_dir,
verbose,
)
print(f"ONNX TF-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}")
print(
f"ONNX TF-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)
def benchmark_onnx(self, inputs):
if self.model_name not in MODELS:
@@ -170,4 +176,6 @@ for currently supported models. Exiting benchmark ONNX."
model_source,
onnx_args,
)
print(f"ONNX ORT-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}")
print(
f"ONNX ORT-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)

16
benchmarks/tests/test_benchmark.py
View File

@@ -38,7 +38,9 @@ class TFHuggingFaceLanguage(tf.Module):
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)
self.m.predict = lambda x, y, z: self.m.call(
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):
@@ -56,7 +58,9 @@ def get_TFhf_model(name):
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(tf.convert_to_tensor(encoded_input[key]), 0)
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0
)
test_input = (
encoded_input["input_ids"],
encoded_input["attention_mask"],
@@ -126,7 +130,9 @@ pytest_benchmark_param = pytest.mark.parametrize(
pytest.param(
False,
"gpu",
marks=pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found"),
marks=pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
),
),
pytest.param(True, "gpu", marks=pytest.mark.skip),
pytest.param(
@@ -155,7 +161,9 @@ pytest_benchmark_param = pytest.mark.parametrize(
)
@pytest_benchmark_param
def test_bench_minilm_torch(dynamic, device):
model, test_input, act_out = get_hf_model("microsoft/MiniLM-L12-H384-uncased")
model, test_input, act_out = get_hf_model(
"microsoft/MiniLM-L12-H384-uncased"
)
shark_module = SharkInference(
model,
(test_input,),

268
generate_sharktank.py
View File

@@ -17,190 +17,142 @@ import csv
import argparse
from shark.shark_importer import SharkImporter
# All generated models and metadata will be saved under this directory.
WORKDIR = os.path.join(os.path.dirname(__file__), "gen_shark_tank")
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
self.upload = upload
print("Setting up for TMP_DIR")
self.workdir = os.path.join(os.path.dirname(__file__), "gen_shark_tank")
print(f"tflite TMP_shark_tank_DIR = {self.workdir}")
os.makedirs(self.workdir, exist_ok=True)
def save_torch_model(torch_model_list):
from tank.model_utils import get_hf_model
from tank.model_utils import get_vision_model
import torch
print("self.torch_model_list: ", self.torch_model_list)
if self.torch_model_list is not None:
self.save_torch_model()
with open(torch_model_list) as csvfile:
torch_reader = csv.reader(csvfile, delimiter=",")
fields = next(torch_reader)
for row in torch_reader:
torch_model_name = row[0]
tracing_required = row[1]
model_type = row[2]
if self.tf_model_list is not None:
self.save_tf_model()
tracing_required = False if tracing_required == "False" else True
print("self.tflite_model_list: ", self.tflite_model_list)
# compile and run tfhub tflite
if self.tflite_model_list is not None:
self.save_tflite_model()
model = None
input = None
if model_type == "vision":
model, input, _ = get_vision_model(torch_model_name)
elif model_type == "hf":
model, input, _ = get_hf_model(torch_model_name)
if self.upload:
print("upload tmp tank to gcp")
os.system("gsutil cp -r ./gen_shark_tank gs://shark_tank/")
torch_model_name = torch_model_name.replace("/", "_")
torch_model_dir = os.path.join(WORKDIR, str(torch_model_name))
os.makedirs(torch_model_dir, exist_ok=True)
def save_torch_model(self):
from tank.model_utils import get_hf_model
from tank.model_utils import get_vision_model, models_dict
import torch
mlir_importer = SharkImporter(
model,
(input,),
frontend="torch",
)
mlir_importer.import_debug(
is_dynamic=False,
tracing_required=tracing_required,
dir=torch_model_dir,
model_name=torch_model_name,
)
with open(self.torch_model_list) as csvfile:
torch_reader = csv.reader(csvfile, delimiter=",")
for row in torch_reader:
torch_model_name = row[0]
print("----------------- torch_model_name", torch_model_name)
is_dynamic = row[1]
tracing_required = row[2]
torch_file = ""
torch_mlir_file = ""
model = []
input = []
if str(torch_model_name)[0:7] == "models.":
print("pretrained model")
model, input, act_out = get_vision_model(models_dict[torch_model_name](pretrained=True))
def save_tf_model(tf_model_list):
print("tf sharktank not implemented yet")
pass
torch_model_name_dir = os.path.join(self.workdir, str(torch_model_name)[7:])
os.makedirs(torch_model_name_dir, exist_ok=True)
print(f"TMP_TORCH_MODELNAME_DIR = {torch_model_name_dir}")
torch_file = "/".join(
[
torch_model_name_dir,
str(torch_model_name)[7:] + ".pt",
]
)
torch_mlir_file = "/".join(
[
torch_model_name_dir,
str(torch_model_name)[7:] + "_torch.mlir",
]
)
else:
model, input, act_out = get_hf_model(str(torch_model_name))
torch_model_name_dir = os.path.join(self.workdir, str(torch_model_name))
os.makedirs(torch_model_name_dir, exist_ok=True)
print(f"TMP_TORCH_MODELNAME_DIR = {torch_model_name_dir}")
loc = torch_model_name.find("/") + 1
torch_model_name = torch_model_name[loc:]
def save_tflite_model(tflite_model_list):
from shark.tflite_utils import TFLitePreprocessor
torch_file = "/".join(
[
torch_model_name_dir,
str(torch_model_name) + ".pt",
]
)
torch_mlir_file = "/".join(
[
torch_model_name_dir,
str(torch_model_name) + "_torch.mlir",
]
)
with open(tflite_model_list) as csvfile:
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(
WORKDIR, str(tflite_model_name)
)
os.makedirs(tflite_model_name_dir, exist_ok=True)
print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")
# save torch model
if os.path.exists(torch_file):
print("Exists", torch_file)
else:
torch.save(model.state_dict(), torch_file)
tflite_tosa_file = "/".join(
[
tflite_model_name_dir,
str(tflite_model_name) + "_tflite.mlir",
]
)
# get mlir model
mlir_importer = SharkImporter(
model,
(input,),
frontend="torch",
)
mlir_model, func_name = mlir_importer.import_mlir(is_dynamic=is_dynamic, tracing_required=tracing_required)
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(str(tflite_model_name))
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
# save mlir model
if os.path.exists(torch_mlir_file):
print("Exists", torch_mlir_file)
else:
mlir_str = mlir_model.operation.get_asm()
with open(torch_mlir_file, "w") as f:
f.write(mlir_str)
print(f"Saved mlir in {torch_mlir_file}")
# Use SharkImporter to get SharkInference input args
my_shark_importer = SharkImporter(
module=tflite_interpreter,
inputs=inputs,
frontend="tflite",
raw_model_file=raw_model_file_path,
)
mlir_model, func_name = my_shark_importer.import_mlir()
print("Torch sharktank not implemented yet")
if os.path.exists(tflite_tosa_file):
print("Exists", tflite_tosa_file)
else:
mlir_str = mlir_model.decode("utf-8")
with open(tflite_tosa_file, "w") as f:
f.write(mlir_str)
print(f"Saved mlir in {tflite_tosa_file}")
def save_tf_model(self):
print("tf sharktank not implemented yet")
def save_tflite_model(self):
from shark.tflite_utils import TFLitePreprocessor
with open(self.tflite_model_list) as csvfile:
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.workdir, str(tflite_model_name))
os.makedirs(tflite_model_name_dir, exist_ok=True)
print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")
tflite_tosa_file = "/".join(
[
tflite_model_name_dir,
str(tflite_model_name) + "_tflite.mlir",
]
)
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(str(tflite_model_name))
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
# Use SharkImporter to get SharkInference input args
my_shark_importer = SharkImporter(
module=tflite_interpreter,
inputs=inputs,
frontend="tflite",
raw_model_file=raw_model_file_path,
)
mlir_model, func_name = my_shark_importer.import_mlir()
if os.path.exists(tflite_tosa_file):
print("Exists", tflite_tosa_file)
else:
mlir_str = mlir_model.decode("utf-8")
with open(tflite_tosa_file, "w") as f:
f.write(mlir_str)
print(f"Saved mlir in {tflite_tosa_file}")
# Validates whether the file is present or not.
def is_valid_file(arg):
if not os.path.exists(arg):
return None
else:
return arg
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--torch_model_list",
type=str,
"--torch_model_csv",
type=lambda x: is_valid_file(x),
default="./tank/pytorch/torch_model_list.csv",
help="""Contains the file with torch_model name and args.
Please see: https://github.com/nod-ai/SHARK/blob/main/tank/pytorch/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",
"--tf_model_csv",
type=lambda x: is_valid_file(x),
default="./tank/tf/tf_model_list.csv",
help="Contains the file with tf model name and args.",
)
parser.add_argument(
"--tflite_model_csv",
type=lambda x: is_valid_file(x),
default="./tank/tflite/tflite_model_list.csv",
help="Contains the file with tf model name and args.",
)
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,
)
if args.torch_model_csv:
save_torch_model(args.torch_model_csv)
if args.tf_model_csv:
save_tf_model(args.torch_model_csv)
if args.tflite_model_csv:
save_tflite_model(args.torch_model_csv)
if args.upload:
print("uploading files to gs://shark_tank/")
os.system("gsutil cp -r ./gen_shark_tank/* gs://shark_tank/")

4
shark/backward_makefx.py
View File

@@ -70,7 +70,9 @@ class MakeFxModule:
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

20
shark/examples/shark_eager/dynamo_demo.py
View File

@@ -8,21 +8,27 @@ 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
def __torch_mlir(fx_graph, *args, **kwargs):
assert isinstance(fx_graph, torch.fx.GraphModule), "Model must be an FX GraphModule."
assert isinstance(
fx_graph, torch.fx.GraphModule
), "Model must be an FX GraphModule."
def _unwrap_single_tuple_return(fx_g: torch.fx.GraphModule):
"""Replace tuple with tuple element in functions that return one-element tuples."""
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],)
@@ -39,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)

14
shark/examples/shark_inference/CLIPModel_tf.py
View File

@@ -18,11 +18,15 @@ class CLIPModule(tf.Module):
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)
self.m.predict = lambda x, y, z: self.m(
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__":
@@ -41,7 +45,11 @@ if __name__ == "__main__":
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()

8
shark/examples/shark_inference/gpt2_tf.py
View File

@@ -30,7 +30,11 @@ if __name__ == "__main__":
text = "I love the distilled version of models."
inputs = tokenizer(text, return_tensors="tf")
shark_module = SharkInference(GPT2Module(), (inputs["input_ids"], inputs["attention_mask"]))
shark_module = SharkInference(
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"])))
print(
shark_module.forward((inputs["input_ids"], inputs["attention_mask"]))
)

12
shark/examples/shark_inference/mhlo_example.py
View File

@@ -13,7 +13,9 @@ arg0 = np.ones((1, 4)).astype(np.float32)
arg1 = np.ones((4, 1)).astype(np.float32)
print("Running shark on cpu backend")
shark_module = SharkInference(mhlo_ir, function_name="forward", device="cpu", mlir_dialect="mhlo")
shark_module = SharkInference(
mhlo_ir, function_name="forward", device="cpu", mlir_dialect="mhlo"
)
# Generate the random inputs and feed into the graph.
x = shark_module.generate_random_inputs()
@@ -21,11 +23,15 @@ shark_module.compile()
print(shark_module.forward(x))
print("Running shark on cuda backend")
shark_module = SharkInference(mhlo_ir, function_name="forward", device="cuda", mlir_dialect="mhlo")
shark_module = SharkInference(
mhlo_ir, function_name="forward", device="cuda", mlir_dialect="mhlo"
)
shark_module.compile()
print(shark_module.forward(x))
print("Running shark on vulkan backend")
shark_module = SharkInference(mhlo_ir, function_name="forward", device="vulkan", mlir_dialect="mhlo")
shark_module = SharkInference(
mhlo_ir, function_name="forward", device="vulkan", mlir_dialect="mhlo"
)
shark_module.compile()
print(shark_module.forward(x))

20
shark/examples/shark_inference/minilm_benchmark_tf.py
View File

@@ -17,10 +17,14 @@ 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)
self.m = TFBertModel.from_pretrained(
"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)
self.m.predict = lambda x, y, z: self.m.call(
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):
@@ -29,7 +33,9 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# Prepping Data
tokenizer = BertTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(
text,
@@ -38,14 +44,18 @@ if __name__ == "__main__":
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(tf.convert_to_tensor(encoded_input[key]), 0)
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"],
)
shark_module = SharkInference(BertModule(), test_input, benchmark_mode=True)
shark_module = SharkInference(
BertModule(), test_input, benchmark_mode=True
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
shark_module.benchmark_all(test_input)

12
shark/examples/shark_inference/minilm_jit.py
View File

@@ -31,11 +31,11 @@ mlir_importer = SharkImporter(
)
# torch hugging face models needs tracing..
(minilm_mlir, func_name), inputs, golden_out = mlir_importer.import_debug(tracing_required=True)
mlir_importer.import_debug(tracing_required=True)
print(golden_out)
# print(golden_out)
shark_module = SharkInference(minilm_mlir, func_name, device="cpu", mlir_dialect="linalg")
shark_module.compile()
result = shark_module.forward((test_input, test_input, test_input))
print("Obtained result", result)
# shark_module = SharkInference(minilm_mlir, func_name, device="cpu", mlir_dialect="linalg")
# shark_module.compile()
# result = shark_module.forward((test_input, test_input, test_input))
# print("Obtained result", result)

16
shark/examples/shark_inference/minilm_tf.py
View File

@@ -17,10 +17,14 @@ 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)
self.m = TFBertModel.from_pretrained(
"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)
self.m.predict = lambda x, y, z: self.m.call(
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):
@@ -29,7 +33,9 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# Prepping Data
tokenizer = BertTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(
text,
@@ -38,7 +44,9 @@ if __name__ == "__main__":
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(tf.convert_to_tensor(encoded_input[key]), 0)
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0
)
shark_module = SharkInference(
BertModule(),

12
shark/examples/shark_inference/resnest.py
View File

@@ -9,7 +9,9 @@ torch.hub.list("zhanghang1989/ResNeSt", force_reload=True)
class ResnestModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = torch.hub.load("zhanghang1989/ResNeSt", "resnest50", pretrained=True)
self.model = torch.hub.load(
"zhanghang1989/ResNeSt", "resnest50", pretrained=True
)
self.model.eval()
def forward(self, input):
@@ -25,11 +27,15 @@ mlir_importer = SharkImporter(
frontend="torch",
)
(vision_mlir, func_name), inputs, golden_out = mlir_importer.import_debug(tracing_required=True)
(vision_mlir, func_name), inputs, golden_out = mlir_importer.import_debug(
tracing_required=True
)
print(golden_out)
shark_module = SharkInference(vision_mlir, func_name, device="cpu", mlir_dialect="linalg")
shark_module = SharkInference(
vision_mlir, func_name, device="cpu", mlir_dialect="linalg"
)
shark_module.compile()
result = shark_module.forward((input))
print("Obtained result", result)

8
shark/examples/shark_inference/resnet50_script.py
View File

@@ -12,14 +12,18 @@ 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"
}
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]),
transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
]
)
img_preprocessed = preprocess(img)

4
shark/examples/shark_training/bert_training.py
View File

@@ -33,7 +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()

13
shark/examples/shark_training/bert_training_load_tf.py
View File

@@ -30,17 +30,24 @@ if __name__ == "__main__":
]
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),
tf.convert_to_tensor(
np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32
),
),
)
shark_module.set_frontend("mhlo")

21
shark/examples/shark_training/bert_training_tf.py
View File

@@ -28,16 +28,22 @@ 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)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m = bert_trainer_model
self.m.predict = lambda x: self.m.call(x, training=False)
self.predict = tf.function(input_signature=[bert_input])(self.m.predict)
self.predict = tf.function(input_signature=[bert_input])(
self.m.predict
)
self.m.learn = lambda x, y: self.m.call(x, training=False)
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@@ -67,13 +73,18 @@ if __name__ == "__main__":
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),
tf.convert_to_tensor(
np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32
),
),
)
shark_module.set_frontend("tensorflow")

4
shark/examples/shark_training/neural_net_training.py
View File

@@ -27,7 +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

16
shark/iree_eager_backend.py
View File

@@ -51,12 +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),
requires_grad=tensor.dtype.type
in {np.float, np.float32, np.float64}
and kwargs.get("requires_grad", False),
)
def compile(self, imported_module: Module):
@@ -66,7 +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,5 +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())

8
shark/iree_utils/benchmark_utils.py
View File

@@ -34,7 +34,9 @@ def tensor_to_type_str(input_tensors: tuple, frontend: str):
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("'", "")
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))
@@ -81,7 +83,9 @@ def run_benchmark_module(benchmark_cl):
Input: benchmark command.
"""
benchmark_path = benchmark_cl[0]
assert os.path.exists(benchmark_path), "Cannot find benchmark_module, Please contact SHARK maintainer on discord."
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)

12
shark/iree_utils/compile_utils.py
View File

@@ -57,7 +57,9 @@ def get_iree_common_args():
]
def compile_module_to_flatbuffer(module, device, frontend, func_name, model_config_path):
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)
@@ -112,7 +114,9 @@ def get_iree_compiled_module(
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)
flatbuffer_blob = compile_module_to_flatbuffer(
module, device, frontend, func_name, model_config_path
)
return get_iree_module(flatbuffer_blob, device, func_name)
@@ -125,7 +129,9 @@ def export_iree_module_to_vmfb(
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)
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}.")

6
shark/iree_utils/cpu_utils.py
View File

@@ -20,7 +20,11 @@ import subprocess
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()
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"

32
shark/iree_utils/gpu_utils.py
View File

@@ -65,24 +65,44 @@ 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:
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.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

22
shark/model_annotation.py
View File

@@ -20,11 +20,15 @@ from typing import List, Dict
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"])
MATMUL_OP_NAMES = set(
["linalg.matmul", "linalg.batch_matmul", "mhlo.dot", "mhlo.dot_general"]
)
idx = 0
def model_annotation(ctx: ir.Context, *, input_contents: str, config_path: str):
def model_annotation(
ctx: ir.Context, *, input_contents: str, config_path: str
):
if os.path.isfile(input_contents):
with open(input_contents, "rb") as f:
input_contents = f.read()
@@ -45,7 +49,9 @@ def model_annotation(ctx: ir.Context, *, input_contents: str, config_path: str):
# More efficient than: print(module)
# - 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))
sys.stdout.buffer.write(
module.operation.get_asm(assume_verified=True, binary=True)
)
return module
@@ -85,7 +91,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:
@@ -149,4 +159,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]
)

4
shark/parser.py
View File

@@ -28,7 +28,9 @@ def dir_file(path):
if os.path.isfile(path):
return path
else:
raise argparse.ArgumentTypeError(f"readable_file:{path} is not a valid file")
raise argparse.ArgumentTypeError(
f"readable_file:{path} is not a valid file"
)
parser = argparse.ArgumentParser(description="SHARK runner.")

20
shark/shark_benchmark_runner.py
View File

@@ -37,10 +37,16 @@ class SharkBenchmarkRunner(SharkRunner):
from_aot: bool = False,
frontend: str = "torch",
):
SharkRunner.__init__(self, model, input, dynamic, device, jit_trace, from_aot, frontend)
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)
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"]:
@@ -144,8 +150,12 @@ class SharkBenchmarkRunner(SharkRunner):
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]
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]

61
shark/shark_downloader.py
View File

@@ -41,7 +41,9 @@ class SharkDownloader:
self.tank_url = tank_url
self.model_type = model_type
self.input_json = input_json # optional if you don't have input
self.input_type = input_type_to_np_dtype[input_type] # optional if you don't have input
self.input_type = input_type_to_np_dtype[
input_type
] # optional if you don't have input
self.mlir_file = None # .mlir file local address.
self.mlir_url = None
self.inputs = None # Input has to be (list of np.array) for sharkInference.forward use
@@ -52,7 +54,9 @@ class SharkDownloader:
print("Error. No tank_url, No model name,Please input either one.")
return
self.workdir = os.path.join(os.path.dirname(__file__), self.local_tank_dir)
self.workdir = os.path.join(
os.path.dirname(__file__), self.local_tank_dir
)
os.makedirs(self.workdir, exist_ok=True)
print(f"TMP_MODEL_DIR = {self.workdir}")
# use model name get dir.
@@ -81,7 +85,9 @@ class SharkDownloader:
def load_json_input(self):
print("load json inputs")
if self.model_type in ["tflite-tosa"]:
input_url = self.tank_url + "/" + str(self.model_name) + "/" + "input.json"
input_url = (
self.tank_url + "/" + str(self.model_name) + "/" + "input.json"
)
input_file = "/".join([self.model_name_dir, str(self.input_json)])
if os.path.exists(input_file):
print("Input has been downloaded before.", input_file)
@@ -92,26 +98,59 @@ class SharkDownloader:
args = []
with open(input_file, "r") as f:
args = json.load(f)
self.inputs = [np.asarray(arg, dtype=self.input_type) for arg in args]
self.inputs = [
np.asarray(arg, dtype=self.input_type) for arg in args
]
else:
print("No json input required for current model type. " "You could call setup_inputs(YOU_INPUTS).")
print(
"No json input required for current model type. "
"You could call setup_inputs(YOU_INPUTS)."
)
return self.inputs
def load_mlir_model(self):
if self.model_type in ["tflite-tosa"]:
self.mlir_url = self.tank_url + "/" + str(self.model_name) + "/" + str(self.model_name) + "_tflite.mlir"
self.mlir_file = "/".join([self.model_name_dir, str(self.model_name) + "_tfite.mlir"])
self.mlir_url = (
self.tank_url
+ "/"
+ str(self.model_name)
+ "/"
+ str(self.model_name)
+ "_tflite.mlir"
)
self.mlir_file = "/".join(
[self.model_name_dir, str(self.model_name) + "_tfite.mlir"]
)
elif self.model_type in ["tensorflow"]:
self.mlir_url = self.tank_url + "/" + str(self.model_name) + "/" + str(self.model_name) + "_tf.mlir"
self.mlir_file = "/".join([self.model_name_dir, str(self.model_name) + "_tf.mlir"])
self.mlir_url = (
self.tank_url
+ "/"
+ str(self.model_name)
+ "/"
+ str(self.model_name)
+ "_tf.mlir"
)
self.mlir_file = "/".join(
[self.model_name_dir, str(self.model_name) + "_tf.mlir"]
)
elif self.model_type in ["torch", "jax", "mhlo", "tosa"]:
self.mlir_url = (
self.tank_url + "/" + str(self.model_name) + "/" + str(self.model_name) + "_" + str(self.model_type) + ".mlir"
self.tank_url
+ "/"
+ str(self.model_name)
+ "/"
+ str(self.model_name)
+ "_"
+ str(self.model_type)
+ ".mlir"
)
self.mlir_file = "/".join(
[
self.model_name_dir,
str(self.model_name) + "_" + str(self.model_type) + ".mlir",
str(self.model_name)
+ "_"
+ str(self.model_type)
+ ".mlir",
]
)
else:

67
shark/shark_importer.py
View File

@@ -2,6 +2,8 @@
"""SHARK Importer"""
import sys
import tempfile
import os
# List of the supported frontends.
supported_frontends = {
@@ -58,7 +60,9 @@ class SharkImporter:
self.inputs = None if len(inputs) == 0 else inputs
self.frontend = frontend
if not self.frontend in supported_frontends:
print(f"The frontend is not in the supported_frontends: {supported_frontends}")
print(
f"The frontend is not in the supported_frontends: {supported_frontends}"
)
sys.exit(1)
self.raw_model_file = raw_model_file
@@ -67,12 +71,16 @@ class SharkImporter:
def _torch_mlir(self, is_dynamic, tracing_required):
from shark.torch_mlir_utils import get_torch_mlir_module
return get_torch_mlir_module(self.module, self.inputs, is_dynamic, tracing_required)
return get_torch_mlir_module(
self.module, self.inputs, is_dynamic, tracing_required
)
def _tf_mlir(self, func_name):
from iree.compiler import tf as tfc
return tfc.compile_module(self.module, exported_names=[func_name], import_only=True)
return tfc.compile_module(
self.module, exported_names=[func_name], import_only=True
)
def _tflite_mlir(self, func_name):
from iree.compiler import tflite as tflitec
@@ -94,7 +102,9 @@ class SharkImporter:
):
if self.frontend in ["torch", "pytorch"]:
if self.inputs == None:
print("Please pass in the inputs, the inputs are required to determine the shape of the mlir_module")
print(
"Please pass in the inputs, the inputs are required to determine the shape of the mlir_module"
)
sys.exit(1)
return self._torch_mlir(is_dynamic, tracing_required), func_name
if self.frontend in ["tf", "tensorflow"]:
@@ -110,25 +120,68 @@ class SharkImporter:
if self.frontend in ["tf", "tensorflow"]:
return [x.numpy() for x in array_tuple]
# Saves `function_name.npy`, `inputs.npz`, `golden_out.npz` and `model_name.mlir` in the directory `dir`.
def save_data(
self, dir, model_name, mlir_data, func_name, inputs, outputs
):
import numpy as np
inputs_name = "inputs.npz"
outputs_name = "golden_out.npz"
func_file_name = "function_name"
model_name_mlir = model_name + ".mlir"
np.savez(os.path.join(dir, inputs_name), *inputs)
np.savez(os.path.join(dir, outputs_name), *outputs)
np.save(os.path.join(dir, func_file_name), np.array(func_name))
mlir_str = mlir_data.operation.get_asm()
with open(os.path.join(dir, model_name_mlir), "w") as mlir_file:
mlir_file.write(mlir_str)
return
def import_debug(
self,
is_dynamic=False,
tracing_required=False,
func_name="forward",
dir=tempfile.gettempdir(),
model_name="model",
):
if self.inputs == None:
print(f"There is no input provided: {self.inputs}, please provide inputs or simply run import_mlir.")
print(
f"There is no input provided: {self.inputs}, please provide inputs or simply run import_mlir."
)
sys.exit(1)
imported_mlir = self.import_mlir(is_dynamic, tracing_required, func_name)
imported_mlir = self.import_mlir(
is_dynamic, tracing_required, func_name
)
# TODO: Make sure that any generic function name is accepted. Currently takes in the default function names.
# TODO: Check for multiple outputs.
if self.frontend in ["torch", "pytorch"]:
import torch
golden_out = self.module(*self.inputs)
if torch.is_tensor(golden_out):
golden_out = tuple(
golden_out.detach().numpy(),
)
else:
golden_out = self.convert_to_numpy(golden_out)
# Save the artifacts in the directory dir.
self.save_data(
dir,
model_name,
imported_mlir[0],
imported_mlir[1],
self.inputs,
golden_out,
)
return (
imported_mlir,
self.convert_to_numpy(self.inputs),
golden_out.detach().numpy(),
golden_out,
)
if self.frontend in ["tf", "tensorflow"]:
golden_out = self.module.forward(*self.inputs)

6
shark/shark_inference.py
View File

@@ -115,5 +115,9 @@ class SharkInference:
shapes, dtype = self._input_info()
inputs = []
for i, j in zip(shapes, dtype):
inputs.append(np.random.uniform(low, high, size=i).astype(dtype_to_np_dtype[j]))
inputs.append(
np.random.uniform(low, high, size=i).astype(
dtype_to_np_dtype[j]
)
)
return tuple(inputs)

9
shark/shark_runner.py
View File

@@ -74,7 +74,10 @@ class SharkRunner:
sys.exit(1)
# Compile the module to get the .vmfb.
(self.iree_compilation_module, self.iree_config,) = get_iree_compiled_module(
(
self.iree_compilation_module,
self.iree_config,
) = get_iree_compiled_module(
self.mlir_module,
self.device,
self.mlir_dialect,
@@ -92,4 +95,6 @@ class SharkRunner:
# 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.mlir_dialect)
return export_iree_module_to_vmfb(
self.model, self.device, dir, self.mlir_dialect
)

8
shark/shark_trainer.py
View File

@@ -69,7 +69,9 @@ 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
@@ -110,7 +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

20
shark/tests/test_shark_importer.py
View File

@@ -24,14 +24,24 @@ def generate_inputs(input_details):
dtype=input_details[0]["dtype"],
)
)
args.append(np.ones(shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]))
args.append(np.zeros(shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]))
args.append(
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -122,7 +132,9 @@ pytest_param = pytest.mark.parametrize(
@pytest_param
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_albert(dynamic, device):
module_tester = AlbertTfliteModuleTester(dynamic=dynamic, device=device)
module_tester.create_and_check_module()

68
shark/tflite_utils.py
View File

@@ -15,7 +15,9 @@ class TFLiteModelUtil:
self.inputs = []
def setup_tflite_interpreter(self):
self.tflite_interpreter = tf.lite.Interpreter(model_path=self.raw_model_file)
self.tflite_interpreter = tf.lite.Interpreter(
model_path=self.raw_model_file
)
self.tflite_interpreter.allocate_tensors()
# default input initialization
return self.get_model_details()
@@ -30,14 +32,20 @@ class TFLiteModelUtil:
self.inputs = inputs
print("invoke_tflite")
for i, input in enumerate(self.inputs):
self.tflite_interpreter.set_tensor(self.input_details[i]["index"], input)
self.tflite_interpreter.set_tensor(
self.input_details[i]["index"], input
)
self.tflite_interpreter.invoke()
# post process tflite_result for compare with mlir_result,
# for tflite the output is a list of numpy.tensor
tflite_results = []
for output_detail in self.output_details:
tflite_results.append(np.array(self.tflite_interpreter.get_tensor(output_detail["index"])))
tflite_results.append(
np.array(
self.tflite_interpreter.get_tensor(output_detail["index"])
)
)
for i in range(len(self.output_details)):
out_dtype = self.output_details[i]["dtype"]
@@ -54,24 +62,40 @@ class TFLitePreprocessor:
model_path=None,
):
self.model_name = model_name
self.input_details = input_details # used for tflite, optional for tf/pytorch
self.output_details = output_details # used for tflite, optional for tf/pytorch
self.input_details = (
input_details # used for tflite, optional for tf/pytorch
)
self.output_details = (
output_details # used for tflite, optional for tf/pytorch
)
self.inputs = []
self.model_path = model_path # url to download the model
self.raw_model_file = None # local address for raw tf/tflite/pytorch model
self.mlir_file = None # local address for .mlir file of tf/tflite/pytorch model
self.raw_model_file = (
None # local address for raw tf/tflite/pytorch model
)
self.mlir_file = (
None # local address for .mlir file of tf/tflite/pytorch model
)
self.mlir_model = None # read of .mlir file
self.output_tensor = None # the raw tf/pytorch/tflite_output_tensor, not mlir_tensor
self.interpreter = None # could be tflite/tf/torch_interpreter in utils
self.output_tensor = (
None # the raw tf/pytorch/tflite_output_tensor, not mlir_tensor
)
self.interpreter = (
None # could be tflite/tf/torch_interpreter in utils
)
self.input_file = None
# 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
print("Setting up for TMP_WORK_DIR")
self.workdir = os.path.join(os.path.dirname(__file__), "./../gen_shark_tank")
self.workdir = os.path.join(
os.path.dirname(__file__), "./../gen_shark_tank"
)
os.makedirs(self.workdir, exist_ok=True)
print(f"TMP_WORK_DIR = {self.workdir}")
@@ -90,13 +114,19 @@ class TFLitePreprocessor:
def load_tflite_model(self):
# use model name get dir.
tflite_model_name_dir = os.path.join(self.workdir, str(self.model_name))
tflite_model_name_dir = os.path.join(
self.workdir, str(self.model_name)
)
os.makedirs(tflite_model_name_dir, exist_ok=True)
print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")
self.raw_model_file = "/".join([tflite_model_name_dir, str(self.model_name) + "_tflite.tflite"])
self.mlir_file = "/".join([tflite_model_name_dir, str(self.model_name) + "_tflite.mlir"])
self.raw_model_file = "/".join(
[tflite_model_name_dir, str(self.model_name) + "_tflite.tflite"]
)
self.mlir_file = "/".join(
[tflite_model_name_dir, str(self.model_name) + "_tflite.mlir"]
)
self.input_file = "/".join([tflite_model_name_dir, "input.json"])
if os.path.exists(self.raw_model_file):
@@ -136,12 +166,18 @@ class TFLitePreprocessor:
self.inputs = []
for tmp_input in input_details:
# print(str(tmp_input["shape"]), tmp_input["dtype"].__name__)
self.inputs.append(np.ones(shape=tmp_input["shape"], dtype=tmp_input["dtype"]))
self.inputs.append(
np.ones(shape=tmp_input["shape"], dtype=tmp_input["dtype"])
)
# save inputs into json file
tmp_json = []
for tmp_input in input_details:
# print(str(tmp_input["shape"]), tmp_input["dtype"].__name__)
tmp_json.append(np.ones(shape=tmp_input["shape"], dtype=tmp_input["dtype"]).tolist())
tmp_json.append(
np.ones(
shape=tmp_input["shape"], dtype=tmp_input["dtype"]
).tolist()
)
with open(self.input_file, "w") as f:
json.dump(tmp_json, f)
return self.inputs

26
shark/torch_mlir_utils.py
View File

@@ -40,7 +40,9 @@ 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
return StringAttr(
module.operation.attributes["torch.debug_module_name"]
).value
def get_input_annotations(inputs: tuple, dynamic: bool) -> list:
@@ -67,7 +69,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,17 +80,27 @@ def shark_jit_trace(module, input: tuple, dynamic: bool, tracing_required: bool)
traced_module = torch.jit.trace_module(module, {"forward": input})
actual_script = traced_module._actual_script_module
export(actual_script.forward)
annotate_args_decorator = annotate_args(get_input_annotations(input, dynamic))
annotate_args_decorator = annotate_args(
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))}
{
"annotations.pkl": pickle.dumps(
extract_serializable_annotations(module)
)
}
)
serializable_test = SerializableTest(
unique_name="", program=torchscript_module_bytes, trace=None
)
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)

4
tank/albert_lite_base/albert_lite_base_tflite_mlir_test.py
View File

@@ -57,7 +57,9 @@ class AlbertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

20
tank/albert_lite_base/albert_lite_base_tflite_test.py
View File

@@ -24,14 +24,24 @@ def generate_inputs(input_details):
dtype=input_details[0]["dtype"],
)
)
args.append(np.ones(shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]))
args.append(np.zeros(shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]))
args.append(
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -126,7 +136,9 @@ class AlbertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

16
tank/arbitrary-image-stylization-v1-256/arbitrary-image-stylization-v1-256_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -43,7 +45,9 @@ class ArbitraryImageStylizationV1TfliteModuleTester:
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
tflite_preprocessor = TFLitePreprocessor(model_name="arbitrary-image-stylization-v1-256")
tflite_preprocessor = TFLitePreprocessor(
model_name="arbitrary-image-stylization-v1-256"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
@@ -84,12 +88,16 @@ class ArbitraryImageStylizationV1TfliteModuleTest(unittest.TestCase):
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = ArbitraryImageStylizationV1TfliteModuleTester(self)
self.module_tester = ArbitraryImageStylizationV1TfliteModuleTester(
self
)
self.module_tester.save_mlir = self.save_mlir
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/birds_V1/birds_V1_tflite_test.py
View File

@@ -31,7 +31,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -123,7 +125,9 @@ class BirdsV1TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/cartoongan/cartoongan_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -88,7 +90,9 @@ class CartoonganTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/deeplabv3/deeplabv3_tflite_test.py
View File

@@ -16,7 +16,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -95,7 +97,9 @@ class DeepLabV3TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/densenet/densenet_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -91,7 +93,9 @@ class DensenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/efficientnet_224_fp32/efficientnet_224_fp32_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class Efficientnet_224_fp32TfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="efficientnet_224_fp32")
tflite_preprocessor = TFLitePreprocessor(
model_name="efficientnet_224_fp32"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class Efficientnet_224_fp32TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/efficientnet_lite0_fp32_2/efficientnet_lite0_fp32_2_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class Efficientnet_lite0_fp32_2TfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="efficientnet_lite0_fp32_2")
tflite_preprocessor = TFLitePreprocessor(
model_name="efficientnet_lite0_fp32_2"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class Efficientnet_lite0_fp32_2TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/efficientnet_lite0_int8_2/efficientnet_lite0_int8_2_tflite_test.py
View File

@@ -24,7 +24,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -56,7 +58,9 @@ class Efficientnet_lite0_int8_2TfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="efficientnet_lite0_int8_2")
tflite_preprocessor = TFLitePreprocessor(
model_name="efficientnet_lite0_int8_2"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -120,7 +124,9 @@ class Efficientnet_lite0_int8_2TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

4
tank/gpt2-64/gpt2-64_tflite_test.py
View File

@@ -25,7 +25,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]

12
tank/inception_v4_299_fp32/inception_v4_299_fp32_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class Inception_v4_299_fp32TfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="inception_v4_299_fp32")
tflite_preprocessor = TFLitePreprocessor(
model_name="inception_v4_299_fp32"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class Inception_v4_299_fp32TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/inception_v4_299_uint8/inception_v4_299_uint8_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -55,7 +57,9 @@ class Inception_v4_299_uint8TfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="inception_v4_299_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="inception_v4_299_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -119,7 +123,9 @@ class Inception_v4_299_uint8TfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/midas/midas_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -91,7 +93,9 @@ class MidasTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/mnasnet_1.0_224/mnasnet_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -91,7 +93,9 @@ class MnasnetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

24
tank/mobilebert-baseline-tf2-float/mobilebert-baseline-tf2-float_tflite_test.py
View File

@@ -15,9 +15,15 @@ def generate_inputs(input_details):
for input in input_details:
print(str(input["shape"]), input["dtype"].__name__)
input_0 = np.asarray(squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"])
input_1 = np.asarray(squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"])
input_2 = np.asarray(squad_data._INPUT_MASK, dtype=input_details[2]["dtype"])
input_0 = np.asarray(
squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"]
)
input_1 = np.asarray(
squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"]
)
input_2 = np.asarray(
squad_data._INPUT_MASK, dtype=input_details[2]["dtype"]
)
return [
input_0.reshape(input_details[0]["shape"]),
input_1.reshape(input_details[1]["shape"]),
@@ -27,7 +33,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -56,7 +64,9 @@ class MobilebertTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilebert-baseline-tf2-float")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilebert-baseline-tf2-float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -120,7 +130,9 @@ class MobilebertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

24
tank/mobilebert-baseline-tf2-quant/mobilebert-baseline-tf2-quant_tflite_test.py
View File

@@ -15,9 +15,15 @@ def generate_inputs(input_details):
for input in input_details:
print(str(input["shape"]), input["dtype"].__name__)
input_0 = np.asarray(squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"])
input_1 = np.asarray(squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"])
input_2 = np.asarray(squad_data._INPUT_MASK, dtype=input_details[2]["dtype"])
input_0 = np.asarray(
squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"]
)
input_1 = np.asarray(
squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"]
)
input_2 = np.asarray(
squad_data._INPUT_MASK, dtype=input_details[2]["dtype"]
)
return [
input_0.reshape(input_details[0]["shape"]),
input_1.reshape(input_details[1]["shape"]),
@@ -27,7 +33,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -56,7 +64,9 @@ class MobilebertTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilebert-baseline-tf2-quant")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilebert-baseline-tf2-quant"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -120,7 +130,9 @@ class MobilebertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

24
tank/mobilebert-edgetpu-s-float/mobilebert-edgetpu-s-float_tflite_test.py
View File

@@ -23,14 +23,24 @@ def generate_inputs(input_details):
dtype=input_details[0]["dtype"],
)
)
args.append(np.ones(shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]))
args.append(np.zeros(shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]))
args.append(
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -61,7 +71,9 @@ class MobilebertTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilebert-edgetpu-s-float")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilebert-edgetpu-s-float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -125,7 +137,9 @@ class MobilebertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

24
tank/mobilebert-edgetpu-s-quant/mobilebert-edgetpu-s-quant_tflite_test.py
View File

@@ -23,14 +23,24 @@ def generate_inputs(input_details):
dtype=input_details[0]["dtype"],
)
)
args.append(np.ones(shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]))
args.append(np.zeros(shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]))
args.append(
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -61,7 +71,9 @@ class MobilebertTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilebert-edgetpu-s-quant")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilebert-edgetpu-s-quant"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -108,7 +120,9 @@ class MobilebertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

20
tank/mobilebert/mobilebert_tflite_test.py
View File

@@ -14,9 +14,15 @@ def generate_inputs(input_details):
for input in input_details:
print(str(input["shape"]), input["dtype"].__name__)
input_0 = np.asarray(squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"])
input_1 = np.asarray(squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"])
input_2 = np.asarray(squad_data._INPUT_MASK, dtype=input_details[2]["dtype"])
input_0 = np.asarray(
squad_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"]
)
input_1 = np.asarray(
squad_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"]
)
input_2 = np.asarray(
squad_data._INPUT_MASK, dtype=input_details[2]["dtype"]
)
return [
input_0.reshape(input_details[0]["shape"]),
input_1.reshape(input_details[1]["shape"]),
@@ -26,7 +32,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -119,7 +127,9 @@ class MobilebertTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v1_224_1.0_float/mobilenet_v1_float_tflite_test.py
View File

@@ -25,7 +25,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -57,7 +59,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v1_224_1.0_float")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v1_224_1.0_float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -121,7 +125,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v1_224_1.0_uint8/mobilenet_v1_uint8_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -55,7 +57,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v1_224_1.0_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v1_224_1.0_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -119,7 +123,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v2_1.00_224_int8/mobilenet_v2_int8_tflite_test.py
View File

@@ -25,7 +25,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -57,7 +59,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v2_1.00_224_int8")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v2_1.00_224_int8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -121,7 +125,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v2_1.0_224/mobilenet_v2_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v2_1.0_224")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v2_1.0_224"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v2_224_1.0_uint8/mobilenet_v2_uint8_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -55,7 +57,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v2_224_1.0_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v2_224_1.0_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -119,7 +123,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v3-large_224_1.0_float/mobilenet_v3_float_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v3-large_224_1.0_float")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v3-large_224_1.0_float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v3-large_224_1.0_uint8/mobilenet_v3_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -55,7 +57,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v3-large_224_1.0_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v3-large_224_1.0_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -119,7 +123,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/mobilenet_v3.5multiavg_1.00_224_int8/mobilenet_v35_int8_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -58,7 +60,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="mobilenet_v3.5multiavg_1.00_224_int8")
tflite_preprocessor = TFLitePreprocessor(
model_name="mobilenet_v3.5multiavg_1.00_224_int8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -122,7 +126,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

15
tank/model_utils.py
View File

@@ -14,9 +14,6 @@ torch.manual_seed(0)
##################### Hugging Face LM Models ###################################
models_dict = {"models.alexnet": models.alexnet}
class HuggingFaceLanguage(torch.nn.Module):
def __init__(self, hf_model_name):
super().__init__()
@@ -44,6 +41,15 @@ def get_hf_model(name):
##################### Torch Vision Models ###################################
vision_models_dict = {
"alexnet": models.alexnet(pretrained=True),
"resnet18": models.resnet18(pretrained=True),
"resnet50": models.resnet50(pretrained=True),
"resnet101": models.resnet101(pretrained=True),
"squeezenet1_0": models.squeezenet1_0(pretrained=True),
"wide_resnet50_2": models.wide_resnet50_2(pretrained=True),
}
class VisionModule(torch.nn.Module):
def __init__(self, model):
@@ -56,8 +62,9 @@ class VisionModule(torch.nn.Module):
def get_vision_model(torch_model):
if isinstance(torch_model, str):
torch_model = vision_models_dict[torch_model]
model = VisionModule(torch_model)
# TODO: Currently the test input is set to (1,128)
test_input = torch.randn(1, 3, 224, 224)
actual_out = model(test_input)
return model, test_input, actual_out

12
tank/model_utils_tf.py
View File

@@ -27,7 +27,9 @@ class TFHuggingFaceLanguage(tf.Module):
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)
self.m.predict = lambda x, y, z: self.m.call(
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,7 +38,9 @@ class TFHuggingFaceLanguage(tf.Module):
def get_TFhf_model(name):
model = TFHuggingFaceLanguage(name)
tokenizer = BertTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(
text,
@@ -45,7 +49,9 @@ def get_TFhf_model(name):
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(tf.convert_to_tensor(encoded_input[key]), 0)
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0
)
test_input = (
encoded_input["input_ids"],
encoded_input["attention_mask"],

12
tank/multi_person_mobilenet_v1_075_float/multi_person_mobilenet_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -41,7 +43,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="multi_person_mobilenet_v1_075_float")
tflite_preprocessor = TFLitePreprocessor(
model_name="multi_person_mobilenet_v1_075_float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -88,7 +92,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

4
tank/nasnet/nasnet_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]

8
tank/person_detect/person_detect_tflite_test.py
View File

@@ -24,14 +24,18 @@ def generate_inputs(input_details):
urllib.request.urlretrieve(img_path, local_path)
shape = input_details[0]["shape"]
im = np.array(Image.open(local_path).resize((shape[1], shape[2]))).astype(input_details[0]["dtype"])
im = np.array(Image.open(local_path).resize((shape[1], shape[2]))).astype(
input_details[0]["dtype"]
)
args = [im.reshape(shape)]
return args
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]

28
tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py
View File

@@ -26,7 +26,9 @@ 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
mlir_importer = SharkImporter(
@@ -34,8 +36,12 @@ class MiniLMModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=True)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -60,25 +66,33 @@ class MiniLMModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

24
tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py
View File

@@ -34,8 +34,12 @@ class AlbertModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=True)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +65,33 @@ class AlbertModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

20
tank/pytorch/alexnet/alexnet_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class AlexnetModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=False)
shark_module = SharkInference(minilm_mlir, func_name, device="cpu", mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=False
)
shark_module = SharkInference(
minilm_mlir, func_name, device="cpu", mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,13 +67,17 @@ class AlexnetModuleTest(unittest.TestCase):
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"

24
tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py
View File

@@ -34,8 +34,12 @@ class BertModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=True)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -60,25 +64,33 @@ class BertModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

24
tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py
View File

@@ -34,8 +34,12 @@ class DistilBertModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=True)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,27 +65,35 @@ class DistilBertModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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.xfail(reason="torch_mlir lowering issues.")
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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.xfail(reason="torch_mlir lowering issues.")
@pytest.mark.skipif(check_device_drivers("vulkan"), reason=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

24
tank/pytorch/mobilebert-uncased/mobilebert-uncased_pytorch_test.py
View File

@@ -34,8 +34,12 @@ class MobileBertUncasedModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=True)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +65,33 @@ class MobileBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module()
@pytest.mark.xfail(reason="golden and original results mismatch")
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

28
tank/pytorch/resnet101/resnet101_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class Resnet101ModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=False)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=False
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +67,33 @@ class Resnet101ModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

28
tank/pytorch/resnet18/resnet18_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class Resnet18ModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=False)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=False
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +67,33 @@ class Resnet18ModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

28
tank/pytorch/resnet50/resnet50_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class Resnet50ModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=False)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=False
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +67,33 @@ class Resnet50ModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

28
tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class SqueezenetModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic, tracing_required=False)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic, tracing_required=False
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +67,33 @@ class SqueezenetModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

7
tank/pytorch/torch_model_blacklist.csv
View File

@@ -1,7 +0,0 @@
distilbert-base-uncased,False,True
models.alexnet,False,False
models.resnet18,False,False
models.resnet50,False,False
models.resnet101,False,False
models.squeezenet1_0,False,False
models.wide_resnet50_2,False,False
1 distilbert-base-uncased False True
2 models.alexnet False False
3 models.resnet18 False False
4 models.resnet50 False False
5 models.resnet101 False False
6 models.squeezenet1_0 False False
7 models.wide_resnet50_2 False False

14
tank/pytorch/torch_model_list.csv
View File

@@ -1,3 +1,11 @@
microsoft/MiniLM-L12-H384-uncased,False,True
albert-base-v2,False,True
bert-base-uncased,False,True
model_name, use_tracing, model_type
microsoft/MiniLM-L12-H384-uncased,True,hf
albert-base-v2,True,hf
bert-base-uncased,True,hf
google/mobilebert-uncased,True,hf
alexnet,False,vision
resnet18,False,vision
resnet50,False,vision
resnet101,False,vision
squeezenet1_0,False,vision
wide_resnet50_2,False,vision
1 microsoft/MiniLM-L12-H384-uncased model_name False use_tracing True model_type
2 albert-base-v2 microsoft/MiniLM-L12-H384-uncased False True True hf
3 bert-base-uncased albert-base-v2 False True True hf
4 bert-base-uncased True hf
5 google/mobilebert-uncased True hf
6 alexnet False vision
7 resnet18 False vision
8 resnet50 False vision
9 resnet101 False vision
10 squeezenet1_0 False vision
11 wide_resnet50_2 False vision

16
tank/pytorch/v_diffusion/cfg_sample.py
View File

@@ -10,7 +10,9 @@ from tqdm import trange
try:
from diffusion import get_model, sampling, utils
except ModuleNotFoundError:
print("You need to download v-diffusion source from https://github.com/crowsonkb/v-diffusion-pytorch")
print(
"You need to download v-diffusion source from https://github.com/crowsonkb/v-diffusion-pytorch"
)
raise
torch.manual_seed(0)
@@ -45,7 +47,9 @@ 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(
@@ -57,7 +61,9 @@ 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)
@@ -85,7 +91,9 @@ def repro(model):
steps = utils.get_spliced_ddpm_cosine_schedule(t)
for i in trange(0, args.n, args.batch_size):
cur_batch_size = min(args.n - i, args.batch_size)
outs = sampling.plms_sample(partial(cfg_model_fn, model), x[i : i + cur_batch_size], steps, {})
outs = sampling.plms_sample(
partial(cfg_model_fn, model), x[i : i + cur_batch_size], steps, {}
)
for j, out in enumerate(outs):
utils.to_pil_image(out).save(f"out_{i + j:05}.png")

28
tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py
View File

@@ -27,7 +27,9 @@ 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
mlir_importer = SharkImporter(
@@ -35,8 +37,12 @@ class WideResnet50ModuleTester:
(input,),
frontend="torch",
)
minilm_mlir, func_name = mlir_importer.import_mlir(is_dynamic=self.dynamic)
shark_module = SharkInference(minilm_mlir, func_name, device=self.device, mlir_dialect="linalg")
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=self.dynamic
)
shark_module = SharkInference(
minilm_mlir, func_name, device=self.device, mlir_dialect="linalg"
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
@@ -61,25 +67,33 @@ class WideResnet50ModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("gpu"))
@pytest.mark.skipif(
check_device_drivers("gpu"), reason=device_driver_info("gpu")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
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=device_driver_info("vulkan"))
@pytest.mark.skipif(
check_device_drivers("vulkan"), reason=device_driver_info("vulkan")
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"

12
tank/resnet_50_224_int8/resnet_50_224_int8_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -44,7 +46,9 @@ class ResnetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="resnet_50_224_int8")
tflite_preprocessor = TFLitePreprocessor(
model_name="resnet_50_224_int8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -91,7 +95,9 @@ class ResnetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/squeezenet/squeezenet_tflite_test.py
View File

@@ -12,7 +12,9 @@ from shark.tflite_utils import TFLitePreprocessor
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -91,7 +93,9 @@ class SequeezeNetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/ssd_mobilenet_v1_320_1.0_float/ssd_mobilenet_v1_float_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -55,7 +57,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="ssd_mobilenet_v1_320_1.0_float")
tflite_preprocessor = TFLitePreprocessor(
model_name="ssd_mobilenet_v1_320_1.0_float"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -119,7 +123,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/ssd_mobilenet_v1_320_1.0_uint8/ssd_mobilenet_v1_uint8_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -52,7 +54,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="ssd_mobilenet_v1_320_1.0_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="ssd_mobilenet_v1_320_1.0_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -116,7 +120,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

16
tank/ssd_mobilenet_v2_face_quant/ssd_mobilenet_v2_face_quant_tflite_test.py
View File

@@ -21,7 +21,9 @@ def generate_inputs(input_details):
workdir = os.path.join(os.path.dirname(__file__), "../tmp", exe_basename)
os.makedirs(workdir, exist_ok=True)
img_path = "https://github.com/google-coral/test_data/raw/master/grace_hopper.bmp"
img_path = (
"https://github.com/google-coral/test_data/raw/master/grace_hopper.bmp"
)
local_path = "/".join([workdir, "grace_hopper.bmp"])
urllib.request.urlretrieve(img_path, local_path)
@@ -33,7 +35,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -62,7 +66,9 @@ class MobilenetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="ssd_mobilenet_v2_face_quant")
tflite_preprocessor = TFLitePreprocessor(
model_name="ssd_mobilenet_v2_face_quant"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -126,7 +132,9 @@ class MobilenetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/ssd_spaghettinet_edgetpu_large/ssd_spaghettinet_edgetpu_large_tflite_test.py
View File

@@ -26,7 +26,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -54,7 +56,9 @@ class SpaghettinetTfliteModuleTester:
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="ssd_spaghettinet_edgetpu_large")
tflite_preprocessor = TFLitePreprocessor(
model_name="ssd_spaghettinet_edgetpu_large"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -118,7 +122,9 @@ class SpaghettinetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

12
tank/ssd_spaghettinet_edgetpu_large_uint8/ssd_spaghettinet_edgetpu_large_uint8_tflite_test.py
View File

@@ -23,7 +23,9 @@ def generate_inputs(input_details):
def compare_results(mlir_results, tflite_results, details):
print("Compare mlir_results VS tflite_results: ")
assert len(mlir_results) == len(tflite_results), "Number of results do not match"
assert len(mlir_results) == len(
tflite_results
), "Number of results do not match"
for i in range(len(details)):
mlir_result = mlir_results[i]
tflite_result = tflite_results[i]
@@ -52,7 +54,9 @@ class SpaghettinetTfliteModuleTester:
shark_args.save_vmfb = self.save_vmfb
# Preprocess to get SharkImporter input args
tflite_preprocessor = TFLitePreprocessor(model_name="ssd_spaghettinet_edgetpu_large_uint8")
tflite_preprocessor = TFLitePreprocessor(
model_name="ssd_spaghettinet_edgetpu_large_uint8"
)
raw_model_file_path = tflite_preprocessor.get_raw_model_file()
inputs = tflite_preprocessor.get_inputs()
tflite_interpreter = tflite_preprocessor.get_interpreter()
@@ -116,7 +120,9 @@ class SpaghettinetTfliteModuleTest(unittest.TestCase):
import sys
@pytest.mark.xfail(sys.platform == "darwin", reason="known macos tflite install issue")
@pytest.mark.xfail(
sys.platform == "darwin", reason="known macos tflite install issue"
)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"

8
tank/tf/automodelmaskedlm.py
View File

@@ -11,7 +11,9 @@ 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)
@@ -39,7 +41,9 @@ supported_models = [
]
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}")

20
tank/tf/bert_large_gen.py
View File

@@ -36,7 +36,9 @@ class BertModule(tf.Module):
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(test_network, num_classes=NUM_CLASSES)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -48,9 +50,15 @@ class BertModule(tf.Module):
@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(
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
]
)
@@ -76,7 +84,9 @@ 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")

18
tank/tf/bert_large_run.py
View File

@@ -38,13 +38,17 @@ class BertModule(tf.Module):
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(test_network, num_classes=NUM_CLASSES)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m = bert_trainer_model
self.m.predict = lambda x: self.m.call(x, training=False)
self.predict = tf.function(input_signature=[bert_input])(self.m.predict)
self.predict = tf.function(input_signature=[bert_input])(
self.m.predict
)
self.m.learn = lambda x, y: self.m.call(x, training=False)
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@@ -71,7 +75,9 @@ 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"
@@ -115,7 +121,11 @@ if __name__ == "__main__":
for i in range(10):
if i == warmup - 1:
start = time.time()
print(BertCompiled.learn(predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))))
print(
BertCompiled.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
end = time.time()
total_time = end - start
print("time: " + str(total_time))

14
tank/tf/bert_large_tf.py
View File

@@ -31,13 +31,17 @@ class BertModule(tf.Module):
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(test_network, num_classes=NUM_CLASSES)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m = bert_trainer_model
self.m.predict = lambda x: self.m.call(x, training=False)
self.predict = tf.function(input_signature=[bert_input])(self.m.predict)
self.predict = tf.function(input_signature=[bert_input])(
self.m.predict
)
self.m.learn = lambda x, y: self.m.call(x, training=False)
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@@ -73,7 +77,11 @@ if __name__ == "__main__":
total_iter = 10
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))))
print(
bert_model.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
if i == warmup - 1:
start = time.time()

24
tank/tf/bert_small_gen.py
View File

@@ -28,10 +28,14 @@ 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)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -43,9 +47,15 @@ class BertModule(tf.Module):
@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(
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
]
)
@@ -71,7 +81,9 @@ 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()

22
tank/tf/bert_small_run.py
View File

@@ -29,16 +29,22 @@ 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)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m = bert_trainer_model
self.m.predict = lambda x: self.m.call(x, training=False)
self.predict = tf.function(input_signature=[bert_input])(self.m.predict)
self.predict = tf.function(input_signature=[bert_input])(
self.m.predict
)
self.m.learn = lambda x, y: self.m.call(x, training=False)
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@@ -65,7 +71,9 @@ 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"
@@ -108,7 +116,11 @@ if __name__ == "__main__":
for i in range(10):
if i == warmup - 1:
start = time.time()
print(BertCompiled.learn(predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))))
print(
BertCompiled.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
end = time.time()
total_time = end - start
print("time: " + str(total_time))

18
tank/tf/bert_small_tf_run.py
View File

@@ -22,16 +22,22 @@ 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)
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m = bert_trainer_model
self.m.predict = lambda x: self.m.call(x, training=False)
self.predict = tf.function(input_signature=[bert_input])(self.m.predict)
self.predict = tf.function(input_signature=[bert_input])(
self.m.predict
)
self.m.learn = lambda x, y: self.m.call(x, training=False)
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@@ -67,7 +73,11 @@ if __name__ == "__main__":
total_iter = 10
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))))
print(
bert_model.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
if i == warmup - 1:
start = time.time()

34
tank/tf/hf_masked_lm/albert-base-v2_test.py
View File

@@ -30,14 +30,20 @@ class AlbertBaseModuleTester:
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if shark_args.save_mlir == True or shark_args.save_vmfb == True or self.save_temps == True:
if (
shark_args.save_mlir == True
or shark_args.save_vmfb == True
or self.save_temps == True
):
repro_path = f"shark_tmp/albert_base_tf_{dynamic}_{device}"
if not os.path.isdir(repro_path):
os.mkdir(repro_path)
shark_args.repro_dir = repro_path
if self.save_temps == True:
temp_dir = tempfile.mkdtemp(prefix="iree_tfs", dir=shark_args.repro_dir)
temp_dir = tempfile.mkdtemp(
prefix="iree_tfs", dir=shark_args.repro_dir
)
np.set_printoptions(threshold=np.inf)
np.save(f"{temp_dir}/input1.npy", input[0])
np.save(f"{temp_dir}/input2.npy", input[1])
@@ -73,7 +79,9 @@ class AlbertBaseModuleTester:
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):
@@ -94,21 +102,29 @@ class AlbertBaseModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skip(reason="Language models currently failing for dynamic case")
@pytest.mark.skip(
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.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")
@pytest.mark.xfail(
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"
@@ -124,7 +140,9 @@ class AlbertBaseModuleTest(unittest.TestCase):
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
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",

38
tank/tf/hf_masked_lm/bert-base-uncased_test.py
View File

@@ -30,14 +30,20 @@ class BertBaseUncasedModuleTester:
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if shark_args.save_mlir == True or shark_args.save_vmfb == True or self.save_temps == True:
if (
shark_args.save_mlir == True
or shark_args.save_vmfb == True
or self.save_temps == True
):
repro_path = f"./shark_tmp/bert_base_uncased_tf_{dynamic}_{device}"
if not os.path.isdir(repro_path):
os.mkdir(repro_path)
shark_args.repro_dir = repro_path
if self.save_temps == True:
temp_dir = tempfile.mkdtemp(prefix="iree_tfs", dir=shark_args.repro_dir)
temp_dir = tempfile.mkdtemp(
prefix="iree_tfs", dir=shark_args.repro_dir
)
np.set_printoptions(threshold=np.inf)
np.save(f"{temp_dir}/input1.npy", input[0])
np.save(f"{temp_dir}/input2.npy", input[1])
@@ -73,7 +79,9 @@ class BertBaseUncasedModuleTester:
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):
@@ -85,28 +93,38 @@ 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"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skip(reason="Language models currently failing for dynamic case")
@pytest.mark.skip(
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")
@pytest.mark.xfail(
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"
@@ -122,7 +140,9 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
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",

38
tank/tf/hf_masked_lm/camembert-base_test.py
View File

@@ -30,14 +30,20 @@ class CamemBertModuleTester:
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if shark_args.save_mlir == True or shark_args.save_vmfb == True or self.save_temps == True:
if (
shark_args.save_mlir == True
or shark_args.save_vmfb == True
or self.save_temps == True
):
repro_path = f"./shark_tmp/camembert_base_tf_{dynamic}_{device}"
if not os.path.isdir(repro_path):
os.mkdir(repro_path)
shark_args.repro_dir = repro_path
if self.save_temps == True:
temp_dir = tempfile.mkdtemp(prefix="iree_tfs", dir=shark_args.repro_dir)
temp_dir = tempfile.mkdtemp(
prefix="iree_tfs", dir=shark_args.repro_dir
)
np.set_printoptions(threshold=np.inf)
np.save(f"{temp_dir}/input1.npy", input[0])
np.save(f"{temp_dir}/input2.npy", input[1])
@@ -73,7 +79,9 @@ class CamemBertModuleTester:
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):
@@ -85,28 +93,38 @@ class CamemBertModuleTest(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.xfail
@pytest.mark.skip(reason="Language models currently failing for dynamic case")
@pytest.mark.skip(
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")
@pytest.mark.xfail(
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"
@@ -122,7 +140,9 @@ class CamemBertModuleTest(unittest.TestCase):
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
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",

42
tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py
View File

@@ -26,18 +26,26 @@ class ConvBertModuleTester:
self.benchmark = benchmark
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("dbmdz/convbert-base-turkish-cased")
model, input, act_out = get_causal_lm_model(
"dbmdz/convbert-base-turkish-cased"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if shark_args.save_mlir == True or shark_args.save_vmfb == True or self.save_temps == True:
if (
shark_args.save_mlir == True
or shark_args.save_vmfb == True
or self.save_temps == True
):
repro_path = f"./shark_tmp/convbert_tf_{dynamic}_{device}"
if not os.path.isdir(repro_path):
os.mkdir(repro_path)
shark_args.repro_dir = repro_path
if self.save_temps == True:
temp_dir = tempfile.mkdtemp(prefix="iree_tfs", dir=shark_args.repro_dir)
temp_dir = tempfile.mkdtemp(
prefix="iree_tfs", dir=shark_args.repro_dir
)
np.set_printoptions(threshold=np.inf)
np.save(f"{temp_dir}/input1.npy", input[0])
np.save(f"{temp_dir}/input2.npy", input[1])
@@ -73,7 +81,9 @@ class ConvBertModuleTester:
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):
@@ -85,28 +95,38 @@ 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"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skip(reason="Language models currently failing for dynamic case")
@pytest.mark.skip(
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")
@pytest.mark.xfail(
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"
@@ -122,7 +142,9 @@ class ConvBertModuleTest(unittest.TestCase):
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
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",

Some files were not shown because too many files have changed in this diff Show More