openvino_executor_runner.cpp can run on several inputs

Author: daniil-lyakhov

examples/openvino/aot/aot_openvino_compiler.py

@@ -5,8 +5,13 @@
 # directory of this source tree for more details.
 
 import argparse
+import os
+import shutil
+import subprocess
+from pathlib import Path
 
 import executorch
+import numpy as np
 import timm
 import torch
 import torchvision.datasets as datasets
@@ -19,9 +24,9 @@
 from sklearn.metrics import accuracy_score
 from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
-from torch.export import ExportedProgram
 from torch.export import export
 from torch.export.exported_program import ExportedProgram
+from torch.fx.passes.graph_drawer import FxGraphDrawer
 from transformers import AutoModel
 
 import nncf
@@ -36,12 +41,14 @@ def load_model(suite: str, model_name: str):
         return timm.create_model(model_name, pretrained=True)
     elif suite == "torchvision":
         if not hasattr(torchvision_models, model_name):
-            raise ValueError(f"Model {model_name} not found in torchvision.")
+            msg = f"Model {model_name} not found in torchvision."
+            raise ValueError(msg)
         return getattr(torchvision_models, model_name)(pretrained=True)
     elif suite == "huggingface":
         return AutoModel.from_pretrained(model_name)
     else:
-        raise ValueError(f"Unsupported model suite: {suite}")
+        msg = f"Unsupported model suite: {suite}"
+        raise ValueError(msg)
 
 
 def load_calibration_dataset(dataset_path: str, suite: str, model: torch.nn.Module):
@@ -61,12 +68,32 @@ def load_calibration_dataset(dataset_path: str, suite: str, model: torch.nn.Modu
     return calibration_dataset
 
 
+def visualize_fx_model(model: torch.fx.GraphModule, output_svg_path: str):
+    g = FxGraphDrawer(model, output_svg_path)
+    g.get_dot_graph().write_svg(output_svg_path)
+
+
+def dump_inputs(calibration_dataset, dest_path):
+    input_files, targets = [], []
+    for idx, data in enumerate(calibration_dataset):
+        feature, target = data
+        targets.append(target)
+        file_name = f"{dest_path}/input_{idx}_0.raw"
+        if not isinstance(feature, torch.Tensor):
+            feature = torch.tensor(feature)
+        feature.detach().numpy().tofile(file_name)
+        input_files.append(file_name)
+
+    return input_files, targets
+
+
 def main(suite: str, model_name: str, input_shape, quantize: bool, dataset_path: str, device: str):
     # Ensure input_shape is a tuple
     if isinstance(input_shape, list):
         input_shape = tuple(input_shape)
     elif not isinstance(input_shape, tuple):
-        raise ValueError("Input shape must be a list or tuple.")
+        msg = "Input shape must be a list or tuple."
+        raise ValueError(msg)
 
     # Load the selected model
     model = load_model(suite, model_name)
@@ -80,11 +107,13 @@ def main(suite: str, model_name: str, input_shape, quantize: bool, dataset_path:
 
     if quantize:
         if suite == "huggingface":
-            raise ValueError("Quantization of {suite} models did not support yet.")
+            msg = f"Quantization of {suite} models did not support yet."
+            raise ValueError(msg)
 
         # Quantize model
         if not dataset_path:
-            raise ValueError("Quantization requires a calibration dataset.")
+            msg = "Quantization requires a calibration dataset."
+            raise ValueError(msg)
         calibration_dataset = load_calibration_dataset(dataset_path, suite, model)
 
         captured_model = aten_dialect.module()
@@ -101,6 +130,7 @@ def transform(x):
             calibration_dataset=nncf.Dataset(calibration_dataset, transform_func=transform),
             fold_quantize=False,
         )
+        visualize_fx_model(quantized_model, f"{model_name}_int8.svg")
 
         aten_dialect: ExportedProgram = export(quantized_model, example_args)
 
@@ -123,54 +153,37 @@ def transform(x):
 
     if quantize:
         print("Start validation of the quantized model:")
-
         # 1: Dump inputs
-        import os
-        import shutil
-
-        dest_path = "tmp_inputs"
-        out_path = "tmp_outputs"
-        targets, input_files = [], []
+        dest_path = Path("tmp_inputs")
+        out_path = Path("tmp_outputs")
         for d in [dest_path, out_path]:
             if os.path.exists(d):
                 shutil.rmtree(d)
             os.makedirs(d)
-        input_list = ""
-        for idx, data in enumerate(calibration_dataset):
-            feature, target = data
-            targets.append(target)
-            file_name = f"{dest_path}/input_{idx}_0.raw"
-            input_list += file_name + " "
-            if not isinstance(feature, torch.Tensor):
-                feature = torch.tensor(feature)
-            feature.detach().numpy().tofile(file_name)
-            input_files.append(file_name)
-
-        inp_list_file = os.path.join(dest_path, "in_list.txt")
+
+        input_files, targets = dump_inputs(calibration_dataset, dest_path)
+        inp_list_file = dest_path / "in_list.txt"
         with open(inp_list_file, "w") as f:
-            input_list = input_list.strip() + "\n"
-            f.write(input_list)
+            f.write("\n".join(input_files) + "\n")
 
         # 2: Run the executor
         print("Run openvino_executor_runner...")
-        import subprocess
 
         subprocess.run(
             [
                 "../../../cmake-openvino-out/examples/openvino/openvino_executor_runner",
                 f"--model_path={model_name}",
                 f"--input_list_path={inp_list_file}",
                 f"--output_folder_path={out_path}",
-                # f"--num_iter={len(input_files)}"
             ]
         )
 
         # 3: load the outputs and compare with the targets
-        import numpy as np
 
         predictions = []
         for i in range(len(input_files)):
-            predictions.append(np.fromfile(os.path.join(out_path, f"output_{i}.raw"), dtype=np.float32))
+            tensor = np.fromfile(out_path / f"output_{i}_0.raw", dtype=np.float32)
+            predictions.append(torch.tensor(np.argmax(tensor)))
 
         acc_top1 = accuracy_score(predictions, targets)
         print(f"acc@1: {acc_top1}")