🔀 Merge changes from #716 (commit: 32cae0b)

Author: shaneahmed

tests/engines/test_patch_predictor.py

@@ -10,10 +10,12 @@
 from typing import TYPE_CHECKING, Callable
 
 import numpy as np
+import torch
 import zarr
 from click.testing import CliRunner
 
-from tiatoolbox import cli
+from tests.conftest import timed
+from tiatoolbox import cli, logger, rcParam
 from tiatoolbox.models import IOPatchPredictorConfig
 from tiatoolbox.models.architecture.vanilla import CNNModel
 from tiatoolbox.models.engine.patch_predictor import PatchPredictor
@@ -555,6 +557,53 @@ def test_engine_run_wsi_annotation_store(
 
     shutil.rmtree(save_dir)
 
+    # ----------------------------------------------------------------------------------
+    # torch.compile
+    # ----------------------------------------------------------------------------------
+    def test_patch_predictor_torch_compile(
+        sample_patch1: Path,
+        sample_patch2: Path,
+        tmp_path: Path,
+    ) -> None:
+        """Test PatchPredictor with torch.compile functionality.
+
+        Args:
+            sample_patch1 (Path): Path to sample patch 1.
+            sample_patch2 (Path): Path to sample patch 2.
+            tmp_path (Path): Path to temporary directory.
+
+        """
+        torch_compile_mode = rcParam["torch_compile_mode"]
+        torch._dynamo.reset()
+        rcParam["torch_compile_mode"] = "default"
+        _, compile_time = timed(
+            test_patch_predictor_api,
+            sample_patch1,
+            sample_patch2,
+            tmp_path,
+        )
+        logger.info("torch.compile default mode: %s", compile_time)
+        torch._dynamo.reset()
+        rcParam["torch_compile_mode"] = "reduce-overhead"
+        _, compile_time = timed(
+            test_patch_predictor_api,
+            sample_patch1,
+            sample_patch2,
+            tmp_path,
+        )
+        logger.info("torch.compile reduce-overhead mode: %s", compile_time)
+        torch._dynamo.reset()
+        rcParam["torch_compile_mode"] = "max-autotune"
+        _, compile_time = timed(
+            test_patch_predictor_api,
+            sample_patch1,
+            sample_patch2,
+            tmp_path,
+        )
+        logger.info("torch.compile max-autotune mode: %s", compile_time)
+        torch._dynamo.reset()
+        rcParam["torch_compile_mode"] = torch_compile_mode
+
 
 # -------------------------------------------------------------------------------------
 # Command Line Interface

tiatoolbox/models/architecture/utils.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import sys
-from typing import Callable, NoReturn
+from typing import NoReturn
 
 import numpy as np
 import torch
@@ -41,7 +41,7 @@ def compile_model(
     model: nn.Module | None = None,
     *,
     mode: str = "default",
-) -> Callable:
+) -> nn.Module:
     """A decorator to compile a model using torch-compile.
 
     Args:
@@ -60,7 +60,7 @@ def compile_model(
               CUDA graphs
 
     Returns:
-        Callable:
+        torch.nn.Module:
             Compiled model.
 
     """
@@ -71,7 +71,7 @@ def compile_model(
     is_torch_compile_compatible()
 
     # This check will be removed when torch.compile is supported in Python 3.12+
-    if sys.version_info >= (3, 12):  # pragma: no cover
+    if sys.version_info > (3, 12):  # pragma: no cover
         logger.warning(
             ("torch-compile is currently not supported in Python 3.12+. ",),
         )

tiatoolbox/models/engine/engine_abc.py

@@ -15,8 +15,9 @@
 from torch import nn
 from typing_extensions import Unpack
 
-from tiatoolbox import DuplicateFilter, logger
+from tiatoolbox import DuplicateFilter, logger, rcParam
 from tiatoolbox.models.architecture import get_pretrained_model
+from tiatoolbox.models.architecture.utils import compile_model
 from tiatoolbox.models.dataset.dataset_abc import PatchDataset, WSIPatchDataset
 from tiatoolbox.models.models_abc import load_torch_model
 from tiatoolbox.utils.misc import (
@@ -355,6 +356,12 @@ def __init__(
             weights=weights,
         )
         self.model.to(device=self.device)
+        self.model = (
+            compile_model(  # for runtime, such as after wrapping with nn.DataParallel
+                model,
+                mode=rcParam["torch_compile_mode"],
+            )
+        )
         self._ioconfig = self.ioconfig  # runtime ioconfig
 
         self.batch_size = batch_size