diff --git a/.github/workflows/test-models.yml b/.github/workflows/test-models.yml new file mode 100644 index 00000000..ef40a1ff --- /dev/null +++ b/.github/workflows/test-models.yml @@ -0,0 +1,33 @@ +# This workflow will install Python dependencies, run tests and lint with a variety of Python versions +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: Validate torch-models on Shark Runtime + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +jobs: + build: + + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9"] + + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v3 + with: + python-version: ${{ matrix.python-version }} + - name: Validate Models + run: | + cd $GITHUB_WORKSPACE + ./setup_venv.sh + source shark.venv/bin/activate + pytest diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000..a2a65783 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +addopts = --verbose -p no:warnings + diff --git a/requirements.txt b/requirements.txt index fd673349..dd500bb2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,4 +9,5 @@ torchvision pybind11 wheel ninja +pytest diff --git a/setup_venv.sh b/setup_venv.sh index 929516ba..9d9df48e 100755 --- a/setup_venv.sh +++ b/setup_venv.sh @@ -44,7 +44,7 @@ fi python -m pip install transformers python -m pip install git+https://github.com/pytorch/functorch.git -python -m pip install . +pip install . --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://github.com/llvm/torch-mlir/releases echo "${Green}Before running examples activate venv with:" echo " ${Green}source $VENV_DIR/bin/activate" diff --git a/shark/examples/test_cpu.py b/shark/examples/test_cpu.py new file mode 100644 index 00000000..f8616f50 --- /dev/null +++ b/shark/examples/test_cpu.py @@ -0,0 +1,118 @@ +from shark.shark_runner import SharkInference + +import torch +import numpy as np +import torchvision.models as models +from transformers import AutoModelForSequenceClassification + +torch.manual_seed(0) + +##################### Hugging Face LM Models ################################### + +class HuggingFaceLanguage(torch.nn.Module): + def __init__(self, hf_model_name): + super().__init__() + self.model = AutoModelForSequenceClassification.from_pretrained( + hf_model_name, # The pretrained model. + num_labels=2, # The number of output labels--2 for binary classification. + output_attentions=False, # Whether the model returns attentions weights. + output_hidden_states=False, # Whether the model returns all hidden-states. + torchscript=True, + ) + + def forward(self, tokens): + return self.model.forward(tokens)[0] + +def get_hf_model(name): + model = HuggingFaceLanguage(name) + # TODO: Currently the test input is set to (1,128) + test_input = torch.randint(2, (1,128)) + actual_out = model(test_input) + return model, test_input, actual_out + + +################################################################################ + + +##################### Torch Vision Models ################################### + +class VisionModule(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.train(False) + + def forward(self, input): + return self.model.forward(input) + + +def get_vision_model(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 + + +################################################################################ + +# Utility function for comparing two tensors. +def compare_tensors(torch_tensor, numpy_tensor): + # setting the absolute and relative tolerance + rtol = 1e-02 + atol = 1e-03 + torch_to_numpy = torch_tensor.detach().numpy() + return np.allclose(torch_to_numpy, numpy_tensor, rtol, atol) + + +def test_bert(): + model, input, act_out = get_hf_model("bert-base-uncased") + shark_module = SharkInference( + model, (input,), device="cpu", jit_trace=True + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results) + + +def test_albert(): + model, input, act_out = get_hf_model("albert-base-v2") + shark_module = SharkInference( + model, (input,), device="cpu", jit_trace=True + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results) + +def test_resnet18(): + model, input, act_out = get_vision_model(models.resnet18(pretrained = True)) + shark_module = SharkInference( + model, + (input,), + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results) + +def test_resnet50(): + model, input, act_out = get_vision_model(models.resnet50(pretrained = True)) + shark_module = SharkInference( + model, + (input,), + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results) + +def test_wide_resnet50(): + model, input, act_out = get_vision_model(models.wide_resnet50_2(pretrained = True)) + shark_module = SharkInference( + model, + (input,), + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results) + +def test_minilm(): + model, input, act_out = get_hf_model("microsoft/MiniLM-L12-H384-uncased") + shark_module = SharkInference( + model, (input,), device="cpu", jit_trace=True + ) + results = shark_module.forward((input,)) + assert True == compare_tensors(act_out, results)