Fix GPU benchmarks on PyTorch model tests. (nod-ai#218)

* Add CUDA_BENCHMARKS option to setup.venv to enable PyTorch benchmarks on CUDA.

* Fix PyTorch GPU benchmarks for tank models.

.github/workflows/test-models.yml

@@ -75,22 +75,22 @@ jobs:
       if: matrix.suite == 'cpu'
       run: |
         cd $GITHUB_WORKSPACE
-        PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
+        PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
         pytest -k 'cpu' --ignore=shark/tests/test_shark_importer.py --ignore=benchmarks/tests/test_hf_benchmark.py --ignore=benchmarks/tests/test_benchmark.py 
 
     - name: Validate GPU Models
       if: matrix.suite == 'gpu'
       run: |
         cd $GITHUB_WORKSPACE
-        PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
+        PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
         pytest -k "gpu" --ignore=shark/tests/test_shark_importer.py --ignore=benchmarks/tests/test_hf_benchmark.py --ignore=benchmarks/tests/test_benchmark.py 
 
     - name: Validate Vulkan Models
       if: matrix.suite == 'vulkan'
       run: |
         cd $GITHUB_WORKSPACE
-        PYTHON=python${{ matrix.python-version }} IMPORTER=1 ./setup_venv.sh
+        PYTHON=python${{ matrix.python-version }} ./setup_venv.sh
         source shark.venv/bin/activate
         pytest -k 'vulkan' --ignore=shark/tests/test_shark_importer.py --ignore=benchmarks/tests/test_hf_benchmark.py --ignore=benchmarks/tests/test_benchmark.py

requirements.txt

@@ -10,3 +10,4 @@ gsutil
 # Testing
 pytest
 pytest-xdist
+Pillow

setup_venv.sh

@@ -98,7 +98,7 @@ if [[ ! -z "${IMPORTER}" ]]; then
   echo "${Yellow}Installing importer tools.."
   if [[ $(uname -s) = 'Linux' ]]; then
     echo "${Yellow}Linux detected.. installing Linux importer tools"
-    $PYTHON -m pip install --upgrade -r "$TD/requirements-importer.txt" -f https://github.com/${RUNTIME}/releases --extra-index-url https://test.pypi.org/simple/ --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    $PYTHON -m pip install --upgrade -r "$TD/requirements-importer.txt" -f https://github.com/${RUNTIME}/releases --extra-index-url https://test.pypi.org/simple/ --extra-index-url https://download.pytorch.org/whl/nightly/cu116
   elif [[ $(uname -s) = 'Darwin' ]]; then
     echo "${Yellow}macOS detected.. installing macOS importer tools"
     #Conda seems to have some problems installing these packages and hope they get resolved upstream.
@@ -108,6 +108,16 @@ fi
 
 $PYTHON -m pip install -e . --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://github.com/llvm/torch-mlir/releases -f https://github.com/${RUNTIME}/releases
 
+if [[ $(uname -s) = 'Linux' ]]; then
+  $PYTHON -m pip uninstall -y torch torchvision
+  $PYTHON -m pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cu116
+  if [ $? -eq 0 ];then
+    echo "Successfully Installed torch + cu116."
+  else
+    echo "Could not install torch + cu116." >&2
+  fi
+fi
+
 if [[ -z "${CONDA_PREFIX}" ]]; then
   echo "${Green}Before running examples activate venv with:"
   echo "  ${Green}source $VENV_DIR/bin/activate"

shark/shark_benchmark_runner.py

@@ -19,6 +19,7 @@
     run_benchmark_module,
 )
 from shark.parser import shark_args
+from tank.model_utils import get_torch_model
 from datetime import datetime
 import time
 import csv
@@ -59,20 +60,33 @@ def setup_cl(self, input_tensors):
             mlir_dialect=self.mlir_dialect,
         )
 
-    def benchmark_frontend(self, inputs):
+    def benchmark_frontend(self, inputs, modelname):
         if self.frontend in ["pytorch", "torch"]:
-            return self.benchmark_torch(inputs)
+            return self.benchmark_torch(modelname)
         elif self.frontend in ["tensorflow", "tf"]:
-            return self.benchmark_tf(inputs)
+            return self.benchmark_tf(inputs, modelname)
+
+    def benchmark_torch(self, modelname):
+        import torch
+
+        if self.device == "gpu":
+            torch.set_default_tensor_type(torch.cuda.FloatTensor)
+        else:
+            torch.set_default_tensor_type(torch.FloatTensor)
+        torch_device = torch.device(
+            "cuda:0" if self.device == "gpu" else "cpu"
+        )
+        HFmodel, input, act_out = get_torch_model(modelname)
+        frontend_model = HFmodel.model
+        frontend_model.to(torch_device)
+        input.to(torch_device)
 
-    def benchmark_torch(self, input_tuple):
-        inputs = input_tuple[0]
         for i in range(shark_args.num_warmup_iterations):
-            self.frontend_model.forward(inputs)
+            frontend_model.forward(input)
 
         begin = time.time()
         for i in range(shark_args.num_iterations):
-            out = self.frontend_model.forward(inputs)
+            out = frontend_model.forward(input)
             if i == shark_args.num_iterations - 1:
                 end = time.time()
                 break
@@ -84,13 +98,13 @@ def benchmark_torch(self, input_tuple):
             f"{((end-begin)/shark_args.num_iterations)*1000}",
         ]
 
-    def benchmark_tf(self, inputs):
+    def benchmark_tf(self, frontend_model, inputs):
         for i in range(shark_args.num_warmup_iterations):
-            self.frontend_model.forward(*inputs)
+            frontend_model.forward(*inputs)
 
         begin = time.time()
         for i in range(shark_args.num_iterations):
-            out = self.frontend_model.forward(*inputs)
+            out = frontend_model.forward(*inputs)
             if i == shark_args.num_iterations - 1:
                 end = time.time()
                 break
@@ -162,12 +176,12 @@ def benchmark_all_csv(
             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, modelname
+                    )[0]
+                    bench_result["ms/iter"] = self.benchmark_frontend(
+                        inputs, modelname
+                    )[1]
                 elif p == "shark_python":
                     bench_result["platform"] = "shark_python"
                     bench_result["iter/sec"] = self.benchmark_python(inputs)[0]

tank/model_utils.py

@@ -2,17 +2,38 @@
 
 import torch
 import numpy as np
+import sys
 
 torch.manual_seed(0)
 
+vision_models = [
+    "alexnet",
+    "resnet101",
+    "resnet18",
+    "resnet50",
+    "squeezenet1_0",
+    "wide_resnet50_2",
+]
+
+
+def get_torch_model(modelname):
+    if modelname in vision_models:
+        return get_vision_model(modelname)
+    else:
+        return get_hf_model(modelname)
+
+
 ##################### Hugging Face LM Models ###################################
 
 
 class HuggingFaceLanguage(torch.nn.Module):
     def __init__(self, hf_model_name):
         super().__init__()
         from transformers import AutoModelForSequenceClassification
+        import transformers as trf
 
+        transformers_path = trf.__path__[0]
+        hf_model_path = f"{transformers_path}/models/{hf_model_name}"
         self.model = AutoModelForSequenceClassification.from_pretrained(
             hf_model_name,  # The pretrained model.
             num_labels=2,  # The number of output labels--2 for binary classification.

tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_test.py

@@ -1,10 +1,8 @@
 from shark.shark_inference import SharkInference
 from shark.iree_utils._common import check_device_drivers, device_driver_info
 from tank.model_utils import compare_tensors
-from shark.parser import shark_args
 from shark.shark_downloader import download_torch_model
 
-import torch
 import unittest
 import numpy as np
 import pytest
@@ -13,20 +11,14 @@
 class MiniLMModuleTester:
     def __init__(
         self,
-        save_mlir=False,
-        save_vmfb=False,
         benchmark=False,
     ):
-        self.save_mlir = save_mlir
-        self.save_vmfb = save_vmfb
         self.benchmark = benchmark
 
     def create_and_check_module(self, dynamic, device):
         model_mlir, func_name, input, act_out = download_torch_model(
             "microsoft/MiniLM-L12-H384-uncased", dynamic
         )
-        shark_args.save_mlir = self.save_mlir
-        shark_args.save_vmfb = self.save_vmfb
 
         # from shark.shark_importer import SharkImporter
         # mlir_importer = SharkImporter(
@@ -50,18 +42,9 @@ def create_and_check_module(self, dynamic, device):
         assert True == compare_tensors(act_out, results)
 
         if self.benchmark == True:
-            import torch
-            from tank.model_utils import get_hf_model
-
-            torch.manual_seed(0)
-
-            model, input, act_out = get_hf_model(
-                "microsoft/MiniLM-L12-H384-uncased"
-            )
-            shark_module.shark_runner.frontend_model = model
             shark_module.shark_runner.benchmark_all_csv(
-                (input,),
-                "MiniLM-L12-H384-uncased",
+                (input),
+                "microsoft/MiniLM-L12-H384-uncased",
                 dynamic,
                 device,
                 "torch",
@@ -72,8 +55,6 @@ class MiniLMModuleTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def configure(self, pytestconfig):
         self.module_tester = MiniLMModuleTester(self)
-        self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
-        self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
         self.module_tester.benchmark = pytestconfig.getoption("benchmark")
 
     def test_module_static_cpu(self):

tank/pytorch/albert-base-v2/albert-base-v2_test.py

@@ -1,10 +1,8 @@
 from shark.shark_inference import SharkInference
 from shark.iree_utils._common import check_device_drivers, device_driver_info
-from tank.model_utils import get_hf_model, compare_tensors
-from shark.parser import shark_args
+from tank.model_utils import compare_tensors
 from shark.shark_downloader import download_torch_model
 
-import torch
 import unittest
 import numpy as np
 import pytest
@@ -13,20 +11,14 @@
 class AlbertModuleTester:
     def __init__(
         self,
-        save_mlir=False,
-        save_vmfb=False,
         benchmark=False,
     ):
-        self.save_mlir = save_mlir
-        self.save_vmfb = save_vmfb
         self.benchmark = benchmark
 
     def create_and_check_module(self, dynamic, device):
         model_mlir, func_name, input, act_out = download_torch_model(
             "albert-base-v2", dynamic
         )
-        shark_args.save_mlir = self.save_mlir
-        shark_args.save_vmfb = self.save_vmfb
 
         # from shark.shark_importer import SharkImporter
         # mlir_importer = SharkImporter(
@@ -50,14 +42,8 @@ def create_and_check_module(self, dynamic, device):
         assert True == compare_tensors(act_out, results)
 
         if self.benchmark == True:
-            import torch
-
-            torch.manual_seed(0)
-
-            model, input, act_out = get_hf_model("albert-base-v2")
-            shark_module.shark_runner.frontend_model = model
             shark_module.shark_runner.benchmark_all_csv(
-                (input,),
+                (input),
                 "albert-base-v2",
                 dynamic,
                 device,
@@ -69,8 +55,6 @@ class AlbertModuleTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def configure(self, pytestconfig):
         self.module_tester = AlbertModuleTester(self)
-        self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
-        self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
         self.module_tester.benchmark = pytestconfig.getoption("benchmark")
 
     def test_module_static_cpu(self):

tank/pytorch/alexnet/alexnet_test.py

@@ -1,7 +1,6 @@
 from shark.shark_inference import SharkInference
 from shark.iree_utils._common import check_device_drivers, device_driver_info
-from tank.model_utils import get_vision_model, compare_tensors
-from shark.parser import shark_args
+from tank.model_utils import compare_tensors
 from shark.shark_downloader import download_torch_model
 
 import unittest
@@ -12,20 +11,14 @@
 class AlexnetModuleTester:
     def __init__(
         self,
-        save_mlir=False,
-        save_vmfb=False,
         benchmark=False,
     ):
-        self.save_mlir = save_mlir
-        self.save_vmfb = save_vmfb
         self.benchmark = benchmark
 
     def create_and_check_module(self, dynamic, device):
         model_mlir, func_name, input, act_out = download_torch_model(
             "alexnet", dynamic
         )
-        shark_args.save_mlir = self.save_mlir
-        shark_args.save_vmfb = self.save_vmfb
 
         # from shark.shark_importer import SharkImporter
         # mlir_importer = SharkImporter(
@@ -49,16 +42,8 @@ def create_and_check_module(self, dynamic, device):
         assert True == compare_tensors(act_out, results)
 
         if self.benchmark == True:
-            import torch
-            import torchvision.models as models
-
-            torch.manual_seed(0)
-            model, input, act_out = get_vision_model(
-                models.alexnet(pretrained=True)
-            )
-            shark_module.shark_runner.frontend_model = model
             shark_module.shark_runner.benchmark_all_csv(
-                (input,),
+                (input),
                 "alexnet",
                 dynamic,
                 device,
@@ -70,8 +55,6 @@ class AlexnetModuleTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def configure(self, pytestconfig):
         self.module_tester = AlexnetModuleTester(self)
-        self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
-        self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
         self.module_tester.benchmark = pytestconfig.getoption("benchmark")
 
     def test_module_static_cpu(self):

tank/pytorch/bench_results.csv

@@ -0,0 +1 @@
+platform,model,dynamic,device,iter/sec,ms/iter,datetime