From 899912e1fffb745deb6d1b74e3cc96db9bab0ec6 Mon Sep 17 00:00:00 2001 From: kkannan Date: Tue, 25 Feb 2025 05:51:57 +0000 Subject: [PATCH] Bringup tt-torch models in forge --- env/core_requirements.txt | 1 + .../models/pytorch/text/albert/test_albert.py | 72 +++++++++++++++++++ .../models/pytorch/text/bloom/__init__.py | 0 .../models/pytorch/text/bloom/test_bloom.py | 47 ++++++++++++ .../models/pytorch/text/bloom/utils/utils.py | 25 +++++++ .../models/pytorch/text/mamba/__init__.py | 0 .../models/pytorch/text/mamba/test_mamba.py | 58 +++++++++++++++ .../pytorch/text/mamba/utils/__init__.py | 0 .../models/pytorch/text/mamba/utils/utils.py | 19 +++++ .../models/pytorch/text/qwen/test_qwen_v2.py | 34 ++++++++- .../models/pytorch/vision/beit/__init__.py | 0 .../models/pytorch/vision/beit/test_beit.py | 39 ++++++++++ .../pytorch/vision/beit/utils/__init__.py | 0 .../models/pytorch/vision/beit/utils/utils.py | 20 ++++++ .../pytorch/vision/glpn_kitti/__int__.py | 0 .../vision/glpn_kitti/test_glpn_kitti.py | 33 +++++++++ .../vision/glpn_kitti/utils/__int__.py | 0 .../pytorch/vision/glpn_kitti/utils/utils.py | 20 ++++++ .../vision/mobilenet/test_mobilenet_v2.py | 27 +++++++ .../pytorch/vision/mobilenet/utils/utils.py | 19 +++++ .../models/pytorch/vision/rmbg/__init__.py | 0 .../models/pytorch/vision/rmbg/test_rmbg.py | 34 +++++++++ .../pytorch/vision/rmbg/utils/__init__.py | 0 .../models/pytorch/vision/rmbg/utils/utils.py | 28 ++++++++ .../models/pytorch/vision/unet/test_unet.py | 27 +++++++ .../pytorch/vision/unet/utils/__init__.py | 0 .../models/pytorch/vision/unet/utils/model.py | 72 +++++++++++++++++++ .../models/pytorch/vision/yolo/test_yolos.py | 33 +++++++++ .../pytorch/vision/yolo/utils/yolos_utils.py | 20 ++++++ forge/test/models/utils.py | 1 + 30 files changed, 628 insertions(+), 1 deletion(-) create mode 100644 forge/test/models/pytorch/text/bloom/__init__.py create mode 100644 forge/test/models/pytorch/text/bloom/test_bloom.py create mode 100644 forge/test/models/pytorch/text/bloom/utils/utils.py create mode 100644 forge/test/models/pytorch/text/mamba/__init__.py create mode 100644 forge/test/models/pytorch/text/mamba/test_mamba.py create mode 100644 forge/test/models/pytorch/text/mamba/utils/__init__.py create mode 100644 forge/test/models/pytorch/text/mamba/utils/utils.py create mode 100644 forge/test/models/pytorch/vision/beit/__init__.py create mode 100644 forge/test/models/pytorch/vision/beit/test_beit.py create mode 100644 forge/test/models/pytorch/vision/beit/utils/__init__.py create mode 100644 forge/test/models/pytorch/vision/beit/utils/utils.py create mode 100644 forge/test/models/pytorch/vision/glpn_kitti/__int__.py create mode 100644 forge/test/models/pytorch/vision/glpn_kitti/test_glpn_kitti.py create mode 100644 forge/test/models/pytorch/vision/glpn_kitti/utils/__int__.py create mode 100644 forge/test/models/pytorch/vision/glpn_kitti/utils/utils.py create mode 100644 forge/test/models/pytorch/vision/rmbg/__init__.py create mode 100644 forge/test/models/pytorch/vision/rmbg/test_rmbg.py create mode 100644 forge/test/models/pytorch/vision/rmbg/utils/__init__.py create mode 100644 forge/test/models/pytorch/vision/rmbg/utils/utils.py create mode 100644 forge/test/models/pytorch/vision/unet/utils/__init__.py create mode 100644 forge/test/models/pytorch/vision/unet/utils/model.py create mode 100644 forge/test/models/pytorch/vision/yolo/test_yolos.py create mode 100644 forge/test/models/pytorch/vision/yolo/utils/yolos_utils.py diff --git a/env/core_requirements.txt b/env/core_requirements.txt index 3144d29f4..d28d028be 100644 --- a/env/core_requirements.txt +++ b/env/core_requirements.txt @@ -51,3 +51,4 @@ pytorch_forecasting==1.0.0 patool openpyxl==3.1.5 GitPython==3.1.44 +kornia==0.8.0 diff --git a/forge/test/models/pytorch/text/albert/test_albert.py b/forge/test/models/pytorch/text/albert/test_albert.py index c1a6d5ee5..90df327a9 100644 --- a/forge/test/models/pytorch/text/albert/test_albert.py +++ b/forge/test/models/pytorch/text/albert/test_albert.py @@ -5,8 +5,11 @@ import torch from transformers import ( AlbertForMaskedLM, + AlbertForQuestionAnswering, + AlbertForSequenceClassification, AlbertForTokenClassification, AlbertTokenizer, + AutoTokenizer, ) import forge @@ -162,3 +165,72 @@ def test_albert_token_classification_pytorch(record_forge_property, size, varian print(f"Context: {sample_text}") print(f"Answer: {predicted_tokens_classes}") + + +@pytest.mark.nightly +def test_albert_question_answering_pytorch(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="albert", + task=Task.QA, + source=Source.HUGGINGFACE, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load Albert tokenizer and model from HuggingFace + tokenizer = download_model(AutoTokenizer.from_pretrained, "twmkn9/albert-base-v2-squad2") + framework_model = download_model( + AlbertForQuestionAnswering.from_pretrained, "twmkn9/albert-base-v2-squad2", return_dict=False + ) + + # Load data sample + question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" + + # Data preprocessing + input_tokens = tokenizer(question, text, return_tensors="pt") + inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]] + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) + + +@pytest.mark.nightly +@pytest.mark.push +def test_albert_sequence_classification_pytorch(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="albert", + task=Task.SEQUENCE_CLASSIFICATION, + source=Source.HUGGINGFACE, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load Albert tokenizer and model from HuggingFace + tokenizer = download_model(AlbertTokenizer.from_pretrained, "textattack/albert-base-v2-imdb") + framework_model = download_model( + AlbertForSequenceClassification.from_pretrained, "textattack/albert-base-v2-imdb", return_dict=False + ) + + # Load data sample + input_text = "Hello, my dog is cute." + + # Data preprocessing + input_tokens = tokenizer(input_text, return_tensors="pt") + inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]] + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/text/bloom/__init__.py b/forge/test/models/pytorch/text/bloom/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/text/bloom/test_bloom.py b/forge/test/models/pytorch/text/bloom/test_bloom.py new file mode 100644 index 000000000..611322181 --- /dev/null +++ b/forge/test/models/pytorch/text/bloom/test_bloom.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +import torch + +import forge +from forge.verify.verify import verify + +from test.models.pytorch.text.bloom.utils.utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + + +# Wrapper to get around past key values +class Wrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, input_ids, attention_mask): + output = self.model(input_ids, None, attention_mask) + return output + + +@pytest.mark.nightly +def test_bloom(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="bloom", + source=Source.HUGGINGFACE, + task=Task.CAUSAL_LM, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + model = load_model() + framework_model = Wrapper(model) + inputs = load_input() + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/text/bloom/utils/utils.py b/forge/test/models/pytorch/text/bloom/utils/utils.py new file mode 100644 index 000000000..01d21674f --- /dev/null +++ b/forge/test/models/pytorch/text/bloom/utils/utils.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def load_model(): + model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b1") + model.config.use_cache = False + model.eval() + return model + + +def load_input(): + test_input = "This is a sample text from " + tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b1", padding_side="left") + inputs = tokenizer.encode_plus( + test_input, + return_tensors="pt", + max_length=32, + padding="max_length", + add_special_tokens=True, + truncation=True, + ) + return [inputs["input_ids"], inputs["attention_mask"]] diff --git a/forge/test/models/pytorch/text/mamba/__init__.py b/forge/test/models/pytorch/text/mamba/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/text/mamba/test_mamba.py b/forge/test/models/pytorch/text/mamba/test_mamba.py new file mode 100644 index 000000000..58462b5cb --- /dev/null +++ b/forge/test/models/pytorch/text/mamba/test_mamba.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +# Reference: https://huggingface.co/state-spaces/mamba-2.8b-hf + +import pytest +import torch + +import forge +from forge.verify.verify import verify + +from test.models.pytorch.text.mamba.utils.utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + + +# Wrapper to return only the output tensor, excluding cache or additional outputs +class Wrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, input_ids): + output = self.model(input_ids) + return output[0] + + +variants = [ + "state-spaces/mamba-790m-hf", + "state-spaces/mamba-2.8b-hf", + "state-spaces/mamba-1.4b-hf", + "state-spaces/mamba-370m-hf", +] + + +@pytest.mark.nightly +@pytest.mark.parametrize("variant", variants) +def test_mamba(record_forge_property, variant): + if variant != "state-spaces/mamba-790m-hf": + pytest.skip("Skipping this variant; only testing the base model (mamba-790m-hf) for now.") + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, model="mamba", variant=variant, task=Task.CAUSAL_LM, source=Source.HUGGINGFACE + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + model = load_model(variant) + framework_model = Wrapper(model) + inputs = load_input(variant) + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/text/mamba/utils/__init__.py b/forge/test/models/pytorch/text/mamba/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/text/mamba/utils/utils.py b/forge/test/models/pytorch/text/mamba/utils/utils.py new file mode 100644 index 000000000..56354a25c --- /dev/null +++ b/forge/test/models/pytorch/text/mamba/utils/utils.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +# Reference: https://huggingface.co/state-spaces/mamba-2.8b-hf + +from transformers import AutoTokenizer, MambaForCausalLM + + +def load_model(variant): + model = MambaForCausalLM.from_pretrained(variant) + model.eval() + return model + + +def load_input(variant): + prompt = "Hey how are you doing?" + tokenizer = AutoTokenizer.from_pretrained(variant) + input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"] + return [input_ids] diff --git a/forge/test/models/pytorch/text/qwen/test_qwen_v2.py b/forge/test/models/pytorch/text/qwen/test_qwen_v2.py index 3a9c4aca8..19f60e2fe 100644 --- a/forge/test/models/pytorch/text/qwen/test_qwen_v2.py +++ b/forge/test/models/pytorch/text/qwen/test_qwen_v2.py @@ -2,7 +2,11 @@ # # SPDX-License-Identifier: Apache-2.0 import pytest -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + Qwen2ForTokenClassification, +) import forge from forge.verify.verify import verify @@ -57,3 +61,31 @@ def test_qwen_clm(record_forge_property, variant): # Model Verification verify(inputs, framework_model, compiled_model) + + +@pytest.mark.nightly +def test_qwen2_token_classification(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, model="qwen_v2", task=Task.TOKEN_CLASSIFICATION, source=Source.HUGGINGFACE + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and tokenizer + framework_model = Qwen2ForTokenClassification.from_pretrained("Qwen/Qwen2-7B") + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B") + + # Prepare input + text = "HuggingFace is a company based in Paris and New York." + model_inputs = tokenizer(text, add_special_tokens=False, return_tensors="pt") + + inputs = [model_inputs["input_ids"], model_inputs["attention_mask"]] + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/beit/__init__.py b/forge/test/models/pytorch/vision/beit/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/beit/test_beit.py b/forge/test/models/pytorch/vision/beit/test_beit.py new file mode 100644 index 000000000..da2b8390b --- /dev/null +++ b/forge/test/models/pytorch/vision/beit/test_beit.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import pytest + +import forge +from forge.verify.verify import verify + +from test.models.pytorch.vision.beit.utils.utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + +variants = ["microsoft/beit-base-patch16-224", "microsoft/beit-large-patch16-224"] + + +@pytest.mark.nightly +@pytest.mark.parametrize("variant", variants) +def test_beit_image_classification(record_forge_property, variant): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="beit", + variant=variant, + source=Source.HUGGINGFACE, + task=Task.IMAGE_CLASSIFICATION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = load_model(variant) + inputs = load_input(variant) + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/beit/utils/__init__.py b/forge/test/models/pytorch/vision/beit/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/beit/utils/utils.py b/forge/test/models/pytorch/vision/beit/utils/utils.py new file mode 100644 index 000000000..fafa01d4d --- /dev/null +++ b/forge/test/models/pytorch/vision/beit/utils/utils.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +import requests +from PIL import Image +from transformers import BeitForImageClassification, BeitImageProcessor + + +def load_model(variant): + model = BeitForImageClassification.from_pretrained(variant) + model.eval() + return model + + +def load_input(variant): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + image = Image.open(requests.get(url, stream=True).raw) + processor = BeitImageProcessor.from_pretrained(variant) + inputs = processor(images=image, return_tensors="pt") + return [inputs["pixel_values"]] diff --git a/forge/test/models/pytorch/vision/glpn_kitti/__int__.py b/forge/test/models/pytorch/vision/glpn_kitti/__int__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/glpn_kitti/test_glpn_kitti.py b/forge/test/models/pytorch/vision/glpn_kitti/test_glpn_kitti.py new file mode 100644 index 000000000..2d15f129c --- /dev/null +++ b/forge/test/models/pytorch/vision/glpn_kitti/test_glpn_kitti.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import forge +from forge.verify.verify import verify + +from test.models.pytorch.vision.glpn_kitti.utils.utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + + +@pytest.mark.nightly +def test_glpn_kitti(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="glpn_kitti", + source=Source.HUGGINGFACE, + task=Task.DEPTH_ESTIMATION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = load_model() + inputs = load_input() + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/glpn_kitti/utils/__int__.py b/forge/test/models/pytorch/vision/glpn_kitti/utils/__int__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/glpn_kitti/utils/utils.py b/forge/test/models/pytorch/vision/glpn_kitti/utils/utils.py new file mode 100644 index 000000000..daac159e0 --- /dev/null +++ b/forge/test/models/pytorch/vision/glpn_kitti/utils/utils.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import requests +from PIL import Image +from transformers import GLPNForDepthEstimation, GLPNImageProcessor + + +def load_model(): + model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti") + model.eval() + return model + + +def load_input(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + image = Image.open(requests.get(url, stream=True).raw) + processor = GLPNImageProcessor.from_pretrained("vinvino02/glpn-kitti") + inputs = processor(images=image, return_tensors="pt") + return [inputs["pixel_values"]] diff --git a/forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v2.py b/forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v2.py index c1a616dc6..ce6656ff7 100644 --- a/forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v2.py +++ b/forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v2.py @@ -21,7 +21,9 @@ from forge.verify.verify import verify from test.models.pytorch.vision.mobilenet.utils.utils import ( + load_input, load_mobilenet_model, + load_model, post_processing, ) from test.models.utils import Framework, Source, Task, build_module_name @@ -283,3 +285,28 @@ def test_mobilenetv2_deeplabv3(record_forge_property, variant): # Model Verification verify(inputs, framework_model, compiled_model) + + +@pytest.mark.nightly +def test_mobilenetv2_torchvision(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="mobilenetv2_torchvision", + source=Source.TORCHVISION, + task=Task.IMAGE_CLASSIFICATION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = load_model() + inputs = load_input() + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/mobilenet/utils/utils.py b/forge/test/models/pytorch/vision/mobilenet/utils/utils.py index 02ece294a..16fdf50af 100644 --- a/forge/test/models/pytorch/vision/mobilenet/utils/utils.py +++ b/forge/test/models/pytorch/vision/mobilenet/utils/utils.py @@ -4,7 +4,9 @@ import os import urllib +import requests import torch +import torchvision.models as models from PIL import Image from torchvision import transforms @@ -59,3 +61,20 @@ def post_processing(output, top_k=5): # Cleanup os.remove("imagenet_classes.txt") os.remove("dog.jpg") + + +def load_model(): + weights = models.MobileNet_V2_Weights.DEFAULT + model = models.mobilenet_v2(weights=weights) + model.eval() + return model + + +def load_input(): + weights = models.MobileNet_V2_Weights.DEFAULT + preprocess = weights.transforms() + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + image = Image.open(requests.get(url, stream=True).raw) + img_t = preprocess(image) + batch_t = torch.unsqueeze(img_t, 0) + return [batch_t] diff --git a/forge/test/models/pytorch/vision/rmbg/__init__.py b/forge/test/models/pytorch/vision/rmbg/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/rmbg/test_rmbg.py b/forge/test/models/pytorch/vision/rmbg/test_rmbg.py new file mode 100644 index 000000000..8e3b99b05 --- /dev/null +++ b/forge/test/models/pytorch/vision/rmbg/test_rmbg.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 + +import forge +from forge.verify.verify import verify + +from test.models.pytorch.vision.rmbg.utils.utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + + +@pytest.mark.nightly +def test_RMBG(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="rmbg_2_0", + source=Source.HUGGINGFACE, + task=Task.IMAGE_SEGMENTATION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = load_model() + inputs = load_input() + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/rmbg/utils/__init__.py b/forge/test/models/pytorch/vision/rmbg/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/rmbg/utils/utils.py b/forge/test/models/pytorch/vision/rmbg/utils/utils.py new file mode 100644 index 000000000..38a834b98 --- /dev/null +++ b/forge/test/models/pytorch/vision/rmbg/utils/utils.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import requests +from PIL import Image +from torchvision import transforms +from transformers import AutoModelForImageSegmentation + + +def load_model(): + model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True) + model.eval() + return model + + +def load_input(): + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + image = Image.open(requests.get(url, stream=True).raw) + image_size = (1024, 1024) + transform_image = transforms.Compose( + [ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + inputs = transform_image(image).unsqueeze(0) + return [inputs] diff --git a/forge/test/models/pytorch/vision/unet/test_unet.py b/forge/test/models/pytorch/vision/unet/test_unet.py index e32565371..692a20175 100644 --- a/forge/test/models/pytorch/vision/unet/test_unet.py +++ b/forge/test/models/pytorch/vision/unet/test_unet.py @@ -23,6 +23,7 @@ import forge from forge.verify.verify import verify +from test.models.pytorch.vision.unet.utils.model import UNET from test.models.utils import Framework, Source, Task, build_module_name from test.utils import download_model @@ -223,3 +224,29 @@ def test_unet_torchhub_pytorch(record_forge_property): # Model Verification verify(inputs, framework_model, compiled_model) + + +# Reference: https://github.com/arief25ramadhan/carvana-unet-segmentation +@pytest.mark.nightly +def test_unet_carvana(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="unet_carvana", + source=Source.GITHUB, + task=Task.IMAGE_SEGMENTATION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = UNET(in_channels=3, out_channels=1) + inputs = [torch.rand((1, 3, 224, 224))] + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/unet/utils/__init__.py b/forge/test/models/pytorch/vision/unet/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/forge/test/models/pytorch/vision/unet/utils/model.py b/forge/test/models/pytorch/vision/unet/utils/model.py new file mode 100644 index 000000000..e4c63182b --- /dev/null +++ b/forge/test/models/pytorch/vision/unet/utils/model.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +# Import libraries +import torch +import torch.nn as nn +from torchvision.transforms import functional + + +class DoubleConv(nn.Module): + def __init__(self, in_channels, out_channels): + super(DoubleConv, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True), + ) + + def forward(self, x): + return self.conv(x) + + +class UNET(nn.Module): + # def __init__(self, in_channels=3, out_channels=1, list_feature_size=[64,128,256,512]): + def __init__(self, in_channels=3, out_channels=1, list_feature_size=[64, 128, 256, 512]): + super(UNET, self).__init__() + self.ups = nn.ModuleList() + self.downs = nn.ModuleList() + self.pool = nn.MaxPool2d(kernel_size=2, stride=2) + + # Down part of unet + for feature_size in list_feature_size: + self.downs.append(DoubleConv(in_channels, feature_size)) + in_channels = feature_size + + # Up part of unet + for feature_size in reversed(list_feature_size): + self.ups.append(nn.ConvTranspose2d(feature_size * 2, feature_size, kernel_size=2, stride=2)) + + self.ups.append(DoubleConv(feature_size * 2, feature_size)) + in_channels = feature_size + + self.bottleneck = DoubleConv(list_feature_size[-1], list_feature_size[-1] * 2) + self.final_conv = nn.Conv2d(list_feature_size[0], out_channels, kernel_size=1) + + def forward(self, x): + list_skip_connections = [] + + for down in self.downs: + x = down(x) + list_skip_connections.append(x) + x = self.pool(x) + + x = self.bottleneck(x) + list_skip_connections = list_skip_connections[::-1] + + concat_skip = torch.empty((1, 1)) + + for idx, up in enumerate(self.ups): + if idx % 2 == 0: + x = up(x) + skip_connection = list_skip_connections[idx // 2] + if x.shape != skip_connection.shape: + x = functional.resize(x, size=skip_connection.shape[2:]) + concat_skip = torch.cat((skip_connection, x), dim=1) + else: + x = up(concat_skip) + + return self.final_conv(x) diff --git a/forge/test/models/pytorch/vision/yolo/test_yolos.py b/forge/test/models/pytorch/vision/yolo/test_yolos.py new file mode 100644 index 000000000..b0e82b6fa --- /dev/null +++ b/forge/test/models/pytorch/vision/yolo/test_yolos.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import forge +from forge.verify.verify import verify + +from test.models.pytorch.vision.yolo.utils.yolos_utils import load_input, load_model +from test.models.utils import Framework, Source, Task, build_module_name + + +@pytest.mark.nightly +def test_yolos(record_forge_property): + + # Build Module Name + module_name = build_module_name( + framework=Framework.PYTORCH, + model="yolos", + source=Source.HUGGINGFACE, + task=Task.OBJECT_DETECTION, + ) + + # Record Forge Property + record_forge_property("model_name", module_name) + + # Load model and input + framework_model = load_model() + inputs = load_input() + + # Forge compile framework model + compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name) + + # Model Verification + verify(inputs, framework_model, compiled_model) diff --git a/forge/test/models/pytorch/vision/yolo/utils/yolos_utils.py b/forge/test/models/pytorch/vision/yolo/utils/yolos_utils.py new file mode 100644 index 000000000..b8c79c78e --- /dev/null +++ b/forge/test/models/pytorch/vision/yolo/utils/yolos_utils.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC + +# SPDX-License-Identifier: Apache-2.0 +import requests +from PIL import Image +from transformers import AutoImageProcessor, AutoModelForObjectDetection + + +def load_model(): + model = AutoModelForObjectDetection.from_pretrained("hustvl/yolos-tiny") + model.eval() + return model + + +def load_input(): + test_input = "http://images.cocodataset.org/val2017/000000039769.jpg" + image = Image.open(requests.get(test_input, stream=True).raw) + image_processor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny") + inputs = image_processor(images=image, return_tensors="pt") + return [inputs["pixel_values"]] diff --git a/forge/test/models/utils.py b/forge/test/models/utils.py index 3789f36a7..46aad79fb 100644 --- a/forge/test/models/utils.py +++ b/forge/test/models/utils.py @@ -36,6 +36,7 @@ class Task(StrEnum): CONDITIONAL_GENERATION = "cond_gen" IMAGE_ENCODING = "img_enc" VISUAL_BACKBONE = "visual_bb" + DEPTH_ESTIMATION = "depth_estimation" class Source(StrEnum):