update compile depth of models & fix UNet issue (#636)

forge/test/model_demos/high_prio/cnn/pytorch/test_blazepose.py

@@ -50,7 +50,7 @@ def test_blazepose_regressor_pytorch(test_device):
     pose_regressor = BlazePoseLandmark()
     pose_regressor.load_weights("mediapipepytorch/blazepose_landmark.pth")
     img2 = [torch.rand(1, 3, 256, 256)]
-    compiled_model = forge.compile(pose_regressor, sample_inputs=[img2], module_name="pt_blazepose_regressor")
+    compiled_model = forge.compile(pose_regressor, sample_inputs=img2, module_name="pt_blazepose_regressor")
 
 
 @pytest.mark.skip(reason="dependent on CCM repo")
@@ -88,4 +88,4 @@ def test_blaze_hand_pytorch(test_device):
     hand_regressor.load_weights("mediapipepytorch/blazehand_landmark.pth")
 
     sample_tensor = [torch.rand(1, 3, 256, 256)]
-    compiled_model = forge.compile(hand_regressor, sample_inputs=[sample_tensor], module_name="pt_hand_regressor")
+    compiled_model = forge.compile(hand_regressor, sample_inputs=sample_tensor, module_name="pt_hand_regressor")

forge/test/model_demos/high_prio/cnn/pytorch/test_clip.py

@@ -103,7 +103,7 @@ def test_clip_pytorch(test_device):
 
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.CONSTEVAL_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
 
     # Load processor and model from HuggingFace
     model_ckpt = "openai/clip-vit-base-patch32"

forge/test/model_demos/high_prio/cnn/pytorch/test_dla.py

@@ -40,7 +40,10 @@
 def test_dla_pytorch(variant, test_device):
     # Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    if variant in ("dla102", "dla102x", "dla102x2", "dla169"):
+        compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
+    else:
+        compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
     func = variants_func[variant]
 
     # Load data sample

forge/test/model_demos/high_prio/cnn/pytorch/test_efficientnet.py

@@ -21,7 +21,7 @@
 
 variants = [
     "efficientnet_b0",
-    # "efficientnet_b4",
+    "efficientnet_b4",
     # "hf_hub:timm/efficientnet_b0.ra_in1k",
     # "hf_hub:timm/efficientnet_b4.ra2_in1k",
     # "hf_hub:timm/efficientnet_b5.in12k_ft_in1k",
@@ -36,7 +36,7 @@ def test_efficientnet_timm(variant, test_device):
 
     # Configuration
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Load model
     framework_model = download_model(timm.create_model, variant, pretrained=True)
@@ -85,7 +85,7 @@ def get_state_dict(self, *args, **kwargs):
 def test_efficientnet_torchvision(variant, test_device):
     # Configuration
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Load model
     if variant == "efficientnet_b0":

forge/test/model_demos/high_prio/cnn/pytorch/test_mobilenet_v1.py

@@ -137,7 +137,7 @@ def forward(self, input):
 def generate_model_mobilenetV1_base_custom_pytorch(test_device, variant):
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Create Forge module from PyTorch model
     model = MobileNetV1(9)
@@ -165,7 +165,7 @@ def test_mobilenetv1_basic(test_device):
 def generate_model_mobilenetv1_imgcls_hf_pytorch(test_device, variant):
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
 
     # Create Forge module from PyTorch model
     preprocessor = download_model(AutoImageProcessor.from_pretrained, variant)
@@ -193,7 +193,7 @@ def test_mobilenetv1_192(test_device):
 def generate_model_mobilenetV1I224_imgcls_hf_pytorch(test_device, variant):
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
 
     # Create Forge module from PyTorch model
     preprocessor = download_model(AutoImageProcessor.from_pretrained, variant)

forge/test/model_demos/high_prio/cnn/pytorch/test_mobilenet_v2.py

@@ -20,7 +20,7 @@
 def generate_model_mobilenetV2_imgcls_torchhub_pytorch(test_device, variant):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
+    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
 
     model = download_model(torch.hub.load, variant, "mobilenet_v2", pretrained=True)
 

forge/test/model_demos/high_prio/cnn/pytorch/test_openpose.py

@@ -326,7 +326,7 @@ def generate_model_openpose_posdet_osmr_pytorch(test_device, variant):
 
     # Configurations
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
 
     # Load model
     framework_model = download_model(ptcv_get_model, variant, pretrained=True)

forge/test/model_demos/high_prio/cnn/pytorch/test_rcnn.py

@@ -57,7 +57,7 @@ def test_rcnn_pytorch(test_device):
 
     # Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.GENERATE_INITIAL_GRAPH
 
     # Proposals generated by selective search were fed to a model in a loop manner to compute features.
     # [Refer line No.151 in https://github.com/object-detection-algorithm/R-CNN/blob/master/py/car_detector.py]

forge/test/model_demos/high_prio/cnn/pytorch/test_resnet.py

@@ -22,7 +22,7 @@ def generate_model_resnet_imgcls_hf_pytorch(variant):
 
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Load data sample
     try:
@@ -47,8 +47,6 @@ def test_resnet(test_device):
         "microsoft/resnet-50",
     )
 
-    compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
     compiled_model = forge.compile(model, sample_inputs=[inputs[0]], module_name="pt_resnet50")
 
 
@@ -60,7 +58,7 @@ def generate_model_resnet_imgcls_timm_pytorch(variant):
 
     # Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Load data sample
     try:

forge/test/model_demos/high_prio/cnn/pytorch/test_resnext.py

@@ -38,7 +38,7 @@ def get_image_tensor():
 def test_resnext_50_torchhub_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(torch.hub.load, "pytorch/vision:v0.10.0", "resnext50_32x4d", pretrained=True)
@@ -55,7 +55,7 @@ def test_resnext_50_torchhub_pytorch(test_device):
 def test_resnext_101_torchhub_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(torch.hub.load, "pytorch/vision:v0.10.0", "resnext101_32x8d", pretrained=True)
@@ -73,7 +73,7 @@ def test_resnext_101_32x8d_fb_wsl_pytorch(test_device):
 
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     # 4 variants
@@ -91,7 +91,7 @@ def test_resnext_101_32x8d_fb_wsl_pytorch(test_device):
 def test_resnext_14_osmr_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(ptcv_get_model, "resnext14_32x4d", pretrained=True)
@@ -109,7 +109,7 @@ def test_resnext_14_osmr_pytorch(test_device):
 def test_resnext_26_osmr_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(ptcv_get_model, "resnext26_32x4d", pretrained=True)
@@ -126,7 +126,7 @@ def test_resnext_26_osmr_pytorch(test_device):
 def test_resnext_50_osmr_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(ptcv_get_model, "resnext50_32x4d", pretrained=True)
@@ -143,7 +143,7 @@ def test_resnext_50_osmr_pytorch(test_device):
 def test_resnext_101_osmr_pytorch(test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     model = download_model(ptcv_get_model, "resnext101_64x4d", pretrained=True)

forge/test/model_demos/high_prio/cnn/pytorch/test_unet.py

@@ -14,7 +14,6 @@
 import pytest
 from pytorchcv.model_provider import get_model as ptcv_get_model
 import segmentation_models_pytorch as smp
-from segmentation_models_pytorch.encoders import get_preprocessing_fn
 
 
 def generate_model_unet_imgseg_osmr_pytorch(variant):
@@ -124,7 +123,7 @@ def generate_model_unet_imgseg_torchhub_pytorch(variant):
 
     model = download_model(
         torch.hub.load,
-        "mateuszforge/brain-segmentation-pytorch",
+        "mateuszbuda/brain-segmentation-pytorch",
         variant,
         in_channels=3,
         out_channels=1,
@@ -135,7 +134,7 @@ def generate_model_unet_imgseg_torchhub_pytorch(variant):
 
     # Download an example input image
     url, filename = (
-        "https://github.com/mateuszforge/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png",
+        "https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png",
         "TCGA_CS_4944.png",
     )
     try:
@@ -156,9 +155,6 @@ def generate_model_unet_imgseg_torchhub_pytorch(variant):
     return model, [img_batch], {}
 
 
-@pytest.mark.skip(
-    reason="Failed to download the model after multiple retries."
-)  # https://github.com/tenstorrent/tt-forge-fe/issues/515
 def test_unet_torchhub_pytorch(test_device):
     model, inputs, _ = generate_model_unet_imgseg_torchhub_pytorch(
         "unet",

forge/test/model_demos/high_prio/cnn/pytorch/test_vgg.py

@@ -27,7 +27,10 @@
 def test_vgg_osmr_pytorch(variant, test_device):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    if variant == "bn_vgg19":
+        compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
+    else:
+        compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
 
     model = download_model(ptcv_get_model, variant, pretrained=True)
     model.eval()
@@ -59,7 +62,7 @@ def test_vgg_19_hf_pytorch(test_device):
 
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.GENERATE_INITIAL_GRAPH
 
     """
     # https://pypi.org/project/vgg-pytorch/
@@ -129,7 +132,7 @@ def test_vgg_bn19_torchhub_pytorch(test_device):
 
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()  # load global compiler config object
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.GENERATE_INITIAL_GRAPH
 
     model = download_model(torch.hub.load, "pytorch/vision:v0.10.0", "vgg19_bn", pretrained=True)
     model.eval()

forge/test/model_demos/high_prio/cnn/pytorch/test_vit.py

@@ -36,8 +36,6 @@ def generate_model_vit_imgcls_hf_pytorch(test_device, variant):
 
 @pytest.mark.parametrize("variant", variants, ids=variants)
 def test_vit_classify_224_hf_pytorch(variant, test_device):
-    compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
     model, inputs, _ = generate_model_vit_imgcls_hf_pytorch(
         test_device,
         variant,

forge/test/model_demos/high_prio/cnn/pytorch/test_vovnet.py

@@ -107,7 +107,6 @@ def test_vovnet_v1_39_stigma_pytorch(test_device, enable_default_dram_parameters
         None,
     )
 
-    compiler_cfg = forge.config._get_global_compiler_config()
     compiled_model = forge.compile(model, sample_inputs=[inputs[0]], module_name=f"pt_vovnet_39_stigma")
 
 

forge/test/model_demos/high_prio/cnn/pytorch/test_yolox.py

@@ -57,7 +57,7 @@ def test_yolox_pytorch(variant, test_device):
 
     # Set PyBuda configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
 
     # prepare model
     weight_name = f"{variant}.pth"

forge/test/model_demos/high_prio/nlp/pytorch/test_llama3.py

@@ -113,7 +113,7 @@ def _update_causal_mask(
 def test_llama3_causal_lm(variant, test_device):
     # Configurations
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.INIT_COMPILE
 
     # Load model (with tokenizer)
     tokenizer = download_model(AutoTokenizer.from_pretrained, variant)

forge/test/model_demos/high_prio/nlp/pytorch/test_phi3.py

@@ -63,7 +63,7 @@ def test_phi3_token_classification(variant, test_device):
 
     # Configurations
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Phi3Config from pretrained variant, disable return_dict and caching.
     config = Phi3Config.from_pretrained(variant)
@@ -96,7 +96,7 @@ def test_phi3_sequence_classification(variant, test_device):
 
     # Configurations
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # Phi3Config from pretrained variant, disable return_dict and caching.
     config = Phi3Config.from_pretrained(variant)

forge/test/model_demos/high_prio/nlp/pytorch/test_t5.py

@@ -82,7 +82,13 @@ def forward(self, decoder_input_ids, encoder_outputs):
 def test_t5_generation(variant, test_device):
 
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = CompileDepth.POST_AUTOGRAD_PASS
+
+    if variant in ["t5-small", "t5-base", "t5-large"]:
+        compiler_cfg.compile_depth = CompileDepth.FINISH_COMPILE
+    elif variant in ["google/flan-t5-small", "google/flan-t5-base"]:
+        compiler_cfg.compile_depth = CompileDepth.SPLIT_GRAPH
+    else:
+        compiler_cfg.compile_depth = CompileDepth.INIT_COMPILE
 
     # Load tokenizer and model from HuggingFace
     # Variants: t5-small, t5-base, t5-large

forge/test/model_demos/high_prio/nlp/pytorch/test_whisper_0.py

@@ -37,7 +37,7 @@
 def generate_model_whisper_congen_hf_pytorch(test_device, variant):
     # Configurations
     compiler_cfg = _get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.POST_AUTOGRAD_PASS
+    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
 
     class Wrapper(torch.nn.Module):
         def __init__(self, model):

forge/test/model_demos/models/wideresnet.py

@@ -17,7 +17,7 @@ def generate_model_wideresnet_imgcls_pytorch(test_device, variant):
 
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     framework_model = download_model(torch.hub.load, "pytorch/vision:v0.10.0", variant, pretrained=True)
@@ -45,7 +45,7 @@ def generate_model_wideresnet_imgcls_timm(test_device, variant):
 
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     framework_model = download_model(timm.create_model, variant, pretrained=True)

forge/test/model_demos/models/xception.py

@@ -14,7 +14,7 @@
 def generate_model_xception_imgcls_timm(test_device, variant):
     # STEP 1: Set Forge configuration parameters
     compiler_cfg = forge.config._get_global_compiler_config()
-    compiler_cfg.compile_depth = forge.CompileDepth.SPLIT_GRAPH
+    compiler_cfg.compile_depth = forge.CompileDepth.FINISH_COMPILE
 
     # STEP 2: Create Forge module from PyTorch model
     framework_model = download_model(timm.create_model, variant, pretrained=True)