⚡️Add torch.compile Functionality

tests/conftest.py

@@ -3,15 +3,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
@@ -568,3 +570,35 @@ def __exit__(self: chdir, *excinfo: object) -> None:
                 os.chdir(self._old_cwd.pop())
 
     return chdir
+
+
+# -------------------------------------------------------------------------------------
+# Utility functions
+# -------------------------------------------------------------------------------------
+
+
+def timed(fn: Callable) -> (Callable, float):
+    """A decorator that times the execution of a function.
+
+    Args:
+        fn (Callable): The function to be timed.
+
+    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()
+        end.record()
+        torch.cuda.synchronize()
+        compile_time = start.elapsed_time(end) / 1000
+    else:
+        start = time.time()
+        result = fn()
+        end = time.time()
+        compile_time = end - start
+    return result, compile_time

tests/test_wsi_registration.py

@@ -5,6 +5,8 @@
 import numpy as np
 import pytest
 
+from tests.conftest import timed
+from tiatoolbox import logger
 from tiatoolbox.tools.registration.wsi_registration import (
     AffineWSITransformer,
     DFBRegister,
@@ -21,10 +23,32 @@
 RNG = np.random.default_rng()  # Numpy Random Generator
 
 
-def test_extract_features(dfbr_features: Path) -> None:
+def test_extract_features_time(dfbr_features: Path, test_count: int = 25) -> None:
+    """Compute time test for CNN based feature extraction function."""
+    compile_time = 0.0
+    eager_compile_time = 0.0
+    for _ in range(test_count):
+        _, _compile_time = timed(lambda: test_extract_features(dfbr_features))
+        compile_time += _compile_time
+    for _ in range(test_count):
+        _, _compile_time = timed(
+            lambda: test_extract_features(dfbr_features, compiled=False),
+        )
+        eager_compile_time += _compile_time
+    compile_time /= test_count
+    eager_compile_time /= test_count
+    logger.info("Time taken for feature extraction (torch.compile): %f", compile_time)
+    logger.info(
+        "Time taken for feature extraction (eager execution): %f",
+        eager_compile_time,
+    )
+    assert compile_time < eager_compile_time
+
+
+def test_extract_features(dfbr_features: Path, *, compiled: bool = True) -> None:
     """Test for CNN based feature extraction function."""
     # dfbr (deep feature based registration).
-    dfbr = DFBRegister()
+    dfbr = DFBRegister(compiled=compiled)
     fixed_img = np.repeat(
         np.expand_dims(
             np.repeat(

tiatoolbox/tools/registration/wsi_registration.py

@@ -328,14 +328,18 @@ class DFBRFeatureExtractor(torch.nn.Module):
 
     """
 
-    def __init__(self: torch.nn.Module) -> None:
+    def __init__(self: torch.nn.Module, *, compiled: bool = True) -> None:
         """Initialize :class:`DFBRFeatureExtractor`."""
         super().__init__()
         output_layers_id: list[str] = ["16", "23", "30"]
         output_layers_key: list[str] = ["block3_pool", "block4_pool", "block5_pool"]
         self.features: dict = dict.fromkeys(output_layers_key, None)
-        self.pretrained: torch.nn.Sequential = torchvision.models.vgg16(
-            weights=VGG16_Weights.IMAGENET1K_V1,
+        self.compiled = compiled
+        self.pretrained: torch.nn.Sequential = torch.compile(
+            torchvision.models.vgg16(
+                weights=VGG16_Weights.IMAGENET1K_V1,
+            ),
+            disable=not compiled,
         ).features
         self.f_hooks = [
             getattr(self.pretrained, layer).register_forward_hook(
@@ -356,6 +360,7 @@ def forward_hook(self: torch.nn.Module, layer_name: str) -> None:
 
         """
 
+        @torch.compile(disable=not self.compiled)
         def hook(
             _module: torch.nn.MaxPool2d,
             _module_input: tuple[torch.Tensor],
@@ -426,11 +431,19 @@ 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),
+        *,
+        compiled: bool = True,
+    ) -> None:
         """Initialize :class:`DFBRegister`."""
         self.patch_size = patch_size
         self.x_scale, self.y_scale = [], []
-        self.feature_extractor = DFBRFeatureExtractor()
+        self.compiled = compiled
+        self.feature_extractor = DFBRFeatureExtractor(
+            compiled=compiled,
+        )
 
     # Make this function private when full pipeline is implemented.
     def extract_features(

tiatoolbox/utils/__init__.py

@@ -9,7 +9,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",