⚡️Add torch.compile Functionality

.github/workflows/mypy-type-check.yml

@@ -6,7 +6,7 @@ on:
   push:
     branches: [ develop, pre-release, master, main ]
   pull_request:
-    branches: [ develop, pre-release, master, main ]
+    branches: [ develop, pre-release, master, main, enhance-torch-compile ]
 
 jobs:
 

.github/workflows/python-package.yml

@@ -8,7 +8,7 @@ on:
     branches: [ develop, pre-release, master, main ]
     tags: v*
   pull_request:
-    branches: [ develop, pre-release, master, main ]
+    branches: [ develop, pre-release, master, main, enhance-torch-compile ]
 
 jobs:
   build:

tests/conftest.py

@@ -4,15 +4,17 @@
 
 import os
 import shutil
+import time
 from pathlib import Path
 from typing import Callable
 
 import pytest
+import torch
 
 import tiatoolbox
 from tiatoolbox import logger
 from tiatoolbox.data import _fetch_remote_sample
-from tiatoolbox.utils.env_detection import running_on_ci
+from tiatoolbox.utils.env_detection import has_gpu, running_on_ci
 
 # -------------------------------------------------------------------------------------
 # Generate Parameterized Tests
@@ -578,3 +580,36 @@ def data_path(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]:
     (tmp_path / "slides").mkdir()
     (tmp_path / "overlays").mkdir()
     return {"base_path": tmp_path}
+
+
+# -------------------------------------------------------------------------------------
+# Utility functions
+# -------------------------------------------------------------------------------------
+
+
+def timed(fn: Callable, *args: object) -> (Callable, float):
+    """A decorator that times the execution of a function.
+
+    Args:
+        fn (Callable): The function to be timed.
+        args (object): Arguments to be passed to the function.
+
+    Returns:
+        A tuple containing the result of the function
+        and the time taken to execute it in seconds.
+    """
+    compile_time = 0.0
+    if has_gpu():
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        result = fn(*args)
+        end.record()
+        torch.cuda.synchronize()
+        compile_time = start.elapsed_time(end) / 1000
+    else:
+        start = time.time()
+        result = fn(*args)
+        end = time.time()
+        compile_time = end - start
+    return result, compile_time

tests/models/test_patch_predictor.py

@@ -13,7 +13,8 @@
 import torch
 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, PatchPredictor
 from tiatoolbox.models.architecture.vanilla import CNNModel
 from tiatoolbox.models.dataset import (
@@ -1226,3 +1227,54 @@ def test_cli_model_multiple_file_mask(remote_sample: Callable, tmp_path: Path) -
     assert tmp_path.joinpath("2.merged.npy").exists()
     assert tmp_path.joinpath("2.raw.json").exists()
     assert tmp_path.joinpath("results.json").exists()
+
+
+# -------------------------------------------------------------------------------------
+# torch.compile
+# -------------------------------------------------------------------------------------
+
+
+def test_patch_predictor_torch_compile(
+    sample_patch1: Path,
+    sample_patch2: Path,
+    tmp_path: Path,
+) -> None:
+    """Test 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_enabled = rcParam["enable_torch_compile"]
+    torch._dynamo.reset()
+    rcParam["enable_torch_compile"] = True
+    # Test torch.compile with default mode
+    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["enable_torch_compile"] = torch_compile_enabled

tests/test_utils.py

@@ -13,6 +13,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import torch
 from PIL import Image
 from requests import HTTPError
 from shapely.geometry import Polygon
@@ -21,6 +22,7 @@
 from tiatoolbox import utils
 from tiatoolbox.annotation.storage import SQLiteStore
 from tiatoolbox.models.architecture import fetch_pretrained_weights
+from tiatoolbox.models.architecture.utils import compile_model
 from tiatoolbox.utils import misc
 from tiatoolbox.utils.exceptions import FileNotSupportedError
 from tiatoolbox.utils.transforms import locsize2bounds
@@ -1819,3 +1821,19 @@ def test_patch_pred_store_persist_ext(tmp_path: pytest.TempPathFactory) -> None:
     # check correct error is raised if coordinates are missing
     with pytest.raises(ValueError, match="coordinates"):
         misc.dict_to_store(patch_output, (1.0, 1.0))
+
+
+def test_torch_compile_already_compiled() -> None:
+    """Test that torch_compile does not recompile a model that is already compiled."""
+    # Create a simple model
+    model = torch.nn.Sequential(torch.nn.Linear(10, 10), torch.nn.Linear(10, 10))
+
+    # Compile the model
+    compiled_model = compile_model(model)
+
+    # Compile the model again
+    recompiled_model = compile_model(compiled_model)
+
+    # Check that the recompiled model
+    # is the same as the original compiled model
+    assert compiled_model == recompiled_model

tiatoolbox/__init__.py

@@ -73,6 +73,8 @@ class _RcParam(TypedDict):
 
     TIATOOLBOX_HOME: Path
     pretrained_model_info: dict[str, dict]
+    enable_torch_compile: bool
+    torch_compile_mode: str
 
 
 def read_registry_files(path_to_registry: str | Path) -> dict:
@@ -102,6 +104,12 @@ def read_registry_files(path_to_registry: str | Path) -> dict:
     "pretrained_model_info": read_registry_files(
         "data/pretrained_model.yaml",
     ),  # Load a dictionary of sample files data (names and urls)
+    "enable_torch_compile": False,
+    # Disable `torch-compile`` by default
+    "torch_compile_mode": "default",
+    # Set ``torch-compile`` mode to ``default`` by default
+    # Options: “default”, “reduce-overhead”, “max-autotune”
+    # or “max-autotune-no-cudagraphs”
 }
 
 

tiatoolbox/models/architecture/__init__.py

@@ -150,6 +150,7 @@ def get_pretrained_model(
     model.load_state_dict(saved_state_dict, strict=True)
 
     # !
+
     io_info = info["ioconfig"]
     creator = locate(f"tiatoolbox.models.engine.{io_info['class']}")
 

tiatoolbox/models/architecture/utils.py

@@ -2,10 +2,60 @@
 
 from __future__ import annotations
 
+import sys
+from typing import Callable
+
 import numpy as np
 import torch
 from torch import nn
 
+from tiatoolbox import logger
+
+
+def compile_model(
+    model: nn.Module | None = None,
+    *,
+    mode: str = "default",
+    disable: bool = False,
+) -> Callable:
+    """A decorator to compile a model using torch-compile.
+
+    Args:
+        model (torch.nn.Module):
+            Model to be compiled.
+        mode (str):
+            Mode to be used for torch-compile. Available modes are
+            `default`, `reduce-overhead`, `max-autotune`, and
+            `max-autotune-no-cudagraphs`.
+        disable (bool):
+            If True, torch-compile will be disabled.
+
+    Returns:
+        Callable:
+            Compiled model.
+
+    """
+    if disable:
+        return model
+
+    # This check will be removed when torch.compile is supported in Python 3.12+
+    if sys.version_info >= (3, 12):  # pragma: no cover
+        logger.warning(
+            ("torch-compile is currently not supported in Python 3.12+. ",),
+        )
+        return model
+
+    if isinstance(
+        model,
+        torch._dynamo.eval_frame.OptimizedModule,  # skipcq: PYL-W0212 # noqa: SLF001
+    ):
+        logger.warning(
+            ("The model is already compiled. ",),
+        )
+        return model
+
+    return torch.compile(model, mode=mode, disable=disable)
+
 
 def centre_crop(
     img: np.ndarray | torch.tensor,

tiatoolbox/models/dataset/classification.py

@@ -163,7 +163,7 @@ class WSIPatchDataset(dataset_abc.PatchDatasetABC):
 
     """
 
-    def __init__(  # skipcq: PY-R1000  # noqa: PLR0913, PLR0915
+    def __init__(  # skipcq: PY-R1000  # noqa: PLR0915 PLR0913
         self: WSIPatchDataset,
         img_path: str | Path,
         mode: str = "wsi",

tiatoolbox/models/engine/multi_task_segmentor.py

@@ -49,7 +49,7 @@
 # Python is yet to be able to natively pickle Object method/static method.
 # Only top-level function is passable to multi-processing as caller.
 # May need 3rd party libraries to use method/static method otherwise.
-def _process_tile_predictions(
+def _process_tile_predictions(  # skipcq: PY-R1000
     ioconfig: IOSegmentorConfig,
     tile_bounds: IntBounds,
     tile_flag: list,

tiatoolbox/models/engine/patch_predictor.py

@@ -11,8 +11,9 @@
 import torch
 import tqdm
 
-from tiatoolbox import logger
+from tiatoolbox import logger, rcParam
 from tiatoolbox.models.architecture import get_pretrained_model
+from tiatoolbox.models.architecture.utils import compile_model
 from tiatoolbox.models.dataset.classification import PatchDataset, WSIPatchDataset
 from tiatoolbox.models.engine.semantic_segmentor import IOSegmentorConfig
 from tiatoolbox.utils import misc, save_as_json
@@ -250,7 +251,13 @@ def __init__(
 
         self.ioconfig = ioconfig  # for storing original
         self._ioconfig = None  # for storing runtime
-        self.model = model  # for runtime, such as after wrapping with nn.DataParallel
+        self.model = (
+            compile_model(  # for runtime, such as after wrapping with nn.DataParallel
+                model,
+                mode=rcParam["torch_compile_mode"],
+                disable=not rcParam["enable_torch_compile"],
+            )
+        )
         self.pretrained_model = pretrained_model
         self.batch_size = batch_size
         self.num_loader_worker = num_loader_workers

tiatoolbox/models/engine/semantic_segmentor.py

@@ -563,7 +563,7 @@ def __init__(
         self.masks = None
 
         self.dataset_class: WSIStreamDataset = dataset_class
-        self.model = model  # original copy
+        self.model = model
         self.pretrained_model = pretrained_model
         self.batch_size = batch_size
         self.num_loader_workers = num_loader_workers

tiatoolbox/models/models_abc.py

@@ -6,8 +6,12 @@
 from typing import TYPE_CHECKING, Any, Callable
 
 import torch
+import torch._dynamo
 from torch import device as torch_device
 
+torch._dynamo.config.suppress_errors = True  # skipcq: PYL-W0212  # noqa: SLF001
+
+
 if TYPE_CHECKING:  # pragma: no cover
     from pathlib import Path
 

tiatoolbox/tools/registration/wsi_registration.py

@@ -431,7 +431,10 @@ class DFBRegister:
 
     """
 
-    def __init__(self: DFBRegister, patch_size: tuple[int, int] = (224, 224)) -> None:
+    def __init__(
+        self: DFBRegister,
+        patch_size: tuple[int, int] = (224, 224),
+    ) -> None:
         """Initialize :class:`DFBRegister`."""
         self.patch_size = patch_size
         self.x_scale: list[float] = []

tiatoolbox/utils/__init__.py

@@ -10,7 +10,14 @@
     visualization,
 )
 
-from .misc import download_data, imread, imwrite, save_as_json, save_yaml, unzip_data
+from .misc import (
+    download_data,
+    imread,
+    imwrite,
+    save_as_json,
+    save_yaml,
+    unzip_data,
+)
 
 __all__ = [
     "imread",