Add changes for e2e inference for passing test cases

forge/test/models/pytorch/multimodal/vilt/test_vilt.py

@@ -56,6 +56,7 @@ def generate_model_vilt_question_answering_hf_pytorch(variant):
 
 
 @pytest.mark.nightly
+@pytest.mark.push
 @pytest.mark.parametrize("variant", variants, ids=variants)
 def test_vilt_question_answering_hf_pytorch(record_forge_property, variant):
     # Build Module Name

forge/test/models/pytorch/vision/dla/test_dla.py

@@ -2,44 +2,29 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import pytest
-import requests
-import torchvision.transforms as transforms
-from PIL import Image
 
 import forge
 from forge.verify.verify import verify
 
-from test.models.pytorch.vision.dla.utils.dla_model import (
-    dla34,
-    dla46_c,
-    dla46x_c,
-    dla60,
-    dla60x,
-    dla60x_c,
-    dla102,
-    dla102x,
-    dla102x2,
-    dla169,
-)
+from test.models.pytorch.vision.dla.utils.utils import load_dla_model, post_processing
 from test.models.utils import Framework, Source, Task, build_module_name
 
-variants_func = {
-    "dla34": dla34,
-    "dla46_c": dla46_c,
-    "dla46x_c": dla46x_c,
-    "dla60x_c": dla60x_c,
-    "dla60": dla60,
-    "dla60x": dla60x,
-    "dla102": dla102,
-    "dla102x": dla102x,
-    "dla102x2": dla102x2,
-    "dla169": dla169,
-}
-variants = list(variants_func.keys())
+variants = [
+    pytest.param("dla34", marks=[pytest.mark.push]),
+    "dla46_c",
+    "dla46x_c",
+    "dla60",
+    "dla60x",
+    "dla60x_c",
+    "dla102",
+    "dla102x",
+    "dla102x2",
+    "dla169",
+]
 
 
 @pytest.mark.nightly
-@pytest.mark.parametrize("variant", variants, ids=variants)
+@pytest.mark.parametrize("variant", variants)
 def test_dla_pytorch(record_forge_property, variant):
     if variant != "dla34":
         pytest.skip("Skipping due to the current CI/CD pipeline limitations")
@@ -52,30 +37,17 @@ def test_dla_pytorch(record_forge_property, variant):
     # Record Forge Property
     record_forge_property("model_name", module_name)
 
-    func = variants_func[variant]
-
-    # Load data sample
-    url = "https://images.rawpixel.com/image_1300/cHJpdmF0ZS9sci9pbWFnZXMvd2Vic2l0ZS8yMDIyLTA1L3BkMTA2LTA0Ny1jaGltXzEuanBn.jpg"
-    image = Image.open(requests.get(url, stream=True).raw)
-
-    # Preprocessing
-    transform = 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]),
-        ]
-    )
-    img_tensor = transform(image).unsqueeze(0)
-
-    framework_model = func(pretrained="imagenet")
-    framework_model.eval()
-
-    inputs = [img_tensor]
+    # Load the model and prepare input data
+    framework_model, inputs = load_dla_model(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)
+
+    # Inference
+    output = compiled_model(*inputs)
+
+    # post processing
+    post_processing(output)

forge/test/models/pytorch/vision/dla/utils/utils.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import os
+import urllib
+
+import requests
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+
+from test.models.pytorch.vision.dla.utils import dla_model
+
+
+def load_dla_model(variant):
+
+    func = getattr(dla_model, variant)
+
+    # Load data sample
+    url = "https://images.rawpixel.com/image_1300/cHJpdmF0ZS9sci9pbWFnZXMvd2Vic2l0ZS8yMDIyLTA1L3BkMTA2LTA0Ny1jaGltXzEuanBn.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+
+    # Preprocessing
+    transform = 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]),
+        ]
+    )
+    img_tensor = transform(image).unsqueeze(0)
+
+    framework_model = func(pretrained="imagenet")
+    framework_model.eval()
+
+    inputs = [img_tensor]
+
+    return framework_model, inputs
+
+
+url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
+
+
+def post_processing(output, top_k=5):
+
+    probabilities = torch.nn.functional.softmax(output[0][0], dim=0)
+    urllib.request.urlretrieve(url, "imagenet_classes.txt")
+
+    with open("imagenet_classes.txt", "r") as f:
+        categories = [s.strip() for s in f.readlines()]
+    topk_prob, topk_catid = torch.topk(probabilities, top_k)
+    for i in range(topk_prob.size(0)):
+        print(categories[topk_catid[i]], topk_prob[i].item())
+
+    # Cleanup
+    os.remove("imagenet_classes.txt")

forge/test/models/pytorch/vision/ghostnet/test_ghostnet.py

@@ -1,23 +1,22 @@
 # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 
 # SPDX-License-Identifier: Apache-2.0
-import urllib
 
 import pytest
-import timm
-from PIL import Image
-from timm.data import resolve_data_config
-from timm.data.transforms_factory import create_transform
 
 import forge
 from forge.verify.verify import verify
 
+from test.models.pytorch.vision.ghostnet.utils.utils import (
+    load_ghostnet_model,
+    post_processing,
+)
 from test.models.utils import Framework, Source, Task, build_module_name
-from test.utils import download_model
 
 variants = ["ghostnet_100"]
 
 
+@pytest.mark.push
 @pytest.mark.nightly
 @pytest.mark.parametrize("variant", variants, ids=variants)
 def test_ghostnet_timm(record_forge_property, variant):
@@ -33,25 +32,17 @@ def test_ghostnet_timm(record_forge_property, variant):
     # Record Forge Property
     record_forge_property("model_name", module_name)
 
-    # STEP 2: Create Forge module from PyTorch model
-    framework_model = download_model(timm.create_model, variant, pretrained=True)
-    framework_model.eval()
-
-    # STEP 3: Prepare input
-    url, filename = (
-        "https://github.com/pytorch/hub/raw/master/images/dog.jpg",
-        "dog.jpg",
-    )
-    urllib.request.urlretrieve(url, filename)
-    img = Image.open(filename)
-    data_config = resolve_data_config({}, model=framework_model)
-    transforms = create_transform(**data_config, is_training=False)
-    img_tensor = transforms(img).unsqueeze(0)
-
-    inputs = [img_tensor]
+    # Load the model and prepare input data
+    framework_model, inputs = load_ghostnet_model(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)
+
+    # Inference
+    output = compiled_model(*inputs)
+
+    # Post processing
+    post_processing(output)

forge/test/models/pytorch/vision/ghostnet/utils/utils.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import os
+import urllib
+
+import timm
+import torch
+from PIL import Image
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
+
+from test.utils import download_model
+
+
+def load_ghostnet_model(variant):
+
+    # Create Forge module from PyTorch model
+    framework_model = download_model(timm.create_model, variant, pretrained=True)
+    framework_model.eval()
+
+    # Prepare input
+    url, filename = (
+        "https://github.com/pytorch/hub/raw/master/images/dog.jpg",
+        "dog.jpg",
+    )
+    urllib.request.urlretrieve(url, filename)
+    img = Image.open(filename)
+    data_config = resolve_data_config({}, model=framework_model)
+    transforms = create_transform(**data_config, is_training=False)
+    img_tensor = transforms(img).unsqueeze(0)
+
+    return framework_model, [img_tensor]
+
+
+url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
+
+
+def post_processing(output, top_k=5):
+
+    probabilities = torch.nn.functional.softmax(output[0][0], dim=0)
+    urllib.request.urlretrieve(url, "imagenet_classes.txt")
+
+    with open("imagenet_classes.txt", "r") as f:
+        categories = [s.strip() for s in f.readlines()]
+    topk_prob, topk_catid = torch.topk(probabilities, top_k)
+    for i in range(topk_prob.size(0)):
+        print(categories[topk_catid[i]], topk_prob[i].item())
+
+    # Cleanup
+    os.remove("imagenet_classes.txt")
+    os.remove("dog.jpg")

forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v1.py

@@ -3,29 +3,22 @@
 # SPDX-License-Identifier: Apache-2.0
 import pytest
 import requests
-import torch
 from PIL import Image
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 import forge
 from forge.verify.verify import verify
 
-from test.models.pytorch.vision.mobilenet.utils.mobilenet_v1 import MobileNetV1
+from test.models.pytorch.vision.mobilenet.utils.utils import (
+    load_mobilenet_model,
+    post_processing,
+)
 from test.models.utils import Framework, Source, Task, build_module_name
 from test.utils import download_model
 
 
-def generate_model_mobilenetV1_base_custom_pytorch():
-    # Create Forge module from PyTorch model
-    model = MobileNetV1(9)
-
-    input_shape = (1, 3, 64, 64)
-    image_tensor = torch.rand(*input_shape)
-
-    return model, [image_tensor], {}
-
-
 @pytest.mark.nightly
+@pytest.mark.push
 def test_mobilenetv1_basic(record_forge_property):
     # Build Module Name
     module_name = build_module_name(
@@ -39,14 +32,21 @@ def test_mobilenetv1_basic(record_forge_property):
     # Record Forge Property
     record_forge_property("model_name", module_name)
 
-    framework_model, inputs, _ = generate_model_mobilenetV1_base_custom_pytorch()
+    # Load the model and prepare input data
+    framework_model, inputs = load_mobilenet_model("mobilenet_v1")
 
     # Forge compile framework model
     compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
 
-    # Model Verification
+    #  Model Verification
     verify(inputs, framework_model, compiled_model)
 
+    # Inference
+    output = compiled_model(*inputs)
+
+    # Post processing
+    post_processing(output)
+
 
 def generate_model_mobilenetv1_imgcls_hf_pytorch(variant):
     # Create Forge module from PyTorch model

forge/test/models/pytorch/vision/mobilenet/test_mobilenet_v2.py

@@ -20,25 +20,16 @@
 import forge
 from forge.verify.verify import verify
 
+from test.models.pytorch.vision.mobilenet.utils.utils import (
+    load_mobilenet_model,
+    post_processing,
+)
 from test.models.utils import Framework, Source, Task, build_module_name
 from test.utils import download_model
 
 
-def generate_model_mobilenetV2_imgcls_torchhub_pytorch():
-    model = download_model(torch.hub.load, "pytorch/vision:v0.10.0", "mobilenet_v2", pretrained=True)
-
-    # Image preprocessing
-    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-    image = Image.open(requests.get(url, stream=True).raw)
-    # TODO : Choose image preprocessor from torchvision,
-    # to make a compatible postprocessing of the predicted class
-    preprocessor = download_model(AutoImageProcessor.from_pretrained, "google/mobilenet_v2_1.0_224")
-    image_tensor = preprocessor(images=image, return_tensors="pt").pixel_values
-
-    return model, [image_tensor], {}
-
-
 @pytest.mark.nightly
+@pytest.mark.push
 def test_mobilenetv2_basic(record_forge_property):
     # Build Module Name
     module_name = build_module_name(
@@ -52,14 +43,21 @@ def test_mobilenetv2_basic(record_forge_property):
     # Record Forge Property
     record_forge_property("model_name", module_name)
 
-    framework_model, inputs, _ = generate_model_mobilenetV2_imgcls_torchhub_pytorch()
+    # Load the model and prepare input data
+    framework_model, inputs = load_mobilenet_model("mobilenet_v2")
 
     # 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)
 
+    # Inference
+    output = compiled_model(*inputs)
+
+    # Post processing
+    post_processing(output)
+
 
 def generate_model_mobilenetV2I96_imgcls_hf_pytorch(variant):
     preprocessor = download_model(AutoImageProcessor.from_pretrained, variant)