🚀 feat(model): Enable Patchcore Training Half Precision (#3047) (#3055)

* 🚀 feat(model): REFACTOR Enable Patchcore Training Half Precision (#3047)

* patchcore infer dtype from model dtype

* lock

* pre-commits

Signed-off-by: Alexander Riedel <alex.riedel@googlemail.com>

---------

Signed-off-by: Alexander Riedel <alex.riedel@googlemail.com>
Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>
Co-authored-by: Ashwin Vaidya <ashwin.vaidya@intel.com>
Co-authored-by: Rajesh Gangireddy <rajesh.gangireddy@intel.com>

---------

Signed-off-by: Alexander Riedel <alex.riedel@googlemail.com>
Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>
Co-authored-by: Alexander Riedel <alex.riedel@googlemail.com>
Co-authored-by: Rajesh Gangireddy <rajesh.gangireddy@intel.com>

Author: 3 people

.github/actions/pytest/action.yaml

@@ -112,7 +112,7 @@ runs:
         uv venv --python "$(command -v python)" .venv
         source .venv/bin/activate
         # Install dependencies with dev extras into the uv-managed environment
-        uv pip install ".[dev]"
+        uv pip install ".[dev,cu124]" # TODO(ashwinvaidya17): See issue #3050
         uv pip install codecov
 
     # Determine which tests to run based on input
@@ -144,6 +144,7 @@ runs:
         INPUTS_MAX_TEST_TIME: ${{ inputs.max-test-time }}
         STEPS_TEST_SCOPE_OUTPUTS_PATH: ${{ steps.test-scope.outputs.path }}
       run: |
+        source .venv/bin/activate
         start_time=$(date +%s)
         set -o pipefail
 
@@ -155,7 +156,7 @@ runs:
         fi
 
         # Run pytest
-        PYTHONPATH=src uv run pytest "$STEPS_TEST_SCOPE_OUTPUTS_PATH" \
+        PYTHONPATH=src pytest "$STEPS_TEST_SCOPE_OUTPUTS_PATH" \
           --numprocesses=0 \
           --durations=10 \
           --durations-min=1.0 \
@@ -225,6 +226,7 @@ runs:
         INPUTS_PYTHON_VERSION: ${{ inputs.python-version }}
 
       run: |
+        source .venv/bin/activate
         uv run codecov --token "$INPUTS_CODECOV_TOKEN" \
                 --file coverage.xml \
                 --flags "$INPUTS_TEST_TYPE_py$INPUTS_PYTHON_VERSION" \

CHANGELOG.md

@@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- 🚀 **model**: Enable Patchcore Training Half Precision by @alexriedel1 in https://github.com/open-edge-platform/anomalib/pull/3047
+
 ### Removed
 
 ### Changed

pyproject.toml

@@ -100,28 +100,32 @@ test = [
 ]
 # PyTorch dependency groups
 cpu = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
     "torchvision>=0.19.0",
 ]
 cu118 = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
     "torchvision>=0.19.0",
 ]
 cu121 = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
     "torchvision>=0.19.0",
 ]
 cu124 = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
+    "torchvision>=0.19.0",
+]
+cu130 = [
+    "torch>=2.4.0", # Note: onnx export fails with torch>=2.9.0 and cuda 13 isn't supported by <=2.8.0
     "torchvision>=0.19.0",
 ]
 rocm = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
     "torchvision>=0.19.0",
     "pytorch-triton-rocm ; sys_platform == 'linux'",
 ]
 xpu = [
-    "torch>=2.4.0",
+    "torch>=2.4.0,<=2.8.0",
     "torchvision>=0.19.0",
     "pytorch-triton-xpu ; sys_platform == 'linux' or sys_platform == 'win32'",
 ]
@@ -158,6 +162,7 @@ conflicts = [
         { extra = "cu118" },
         { extra = "cu121" },
         { extra = "cu124" },
+        { extra = "cu130" },
         { extra = "rocm" },
         { extra = "xpu" },
     ],
@@ -184,6 +189,11 @@ name = "pytorch-cu124"
 url = "https://download.pytorch.org/whl/cu124"
 explicit = true
 
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
 [[tool.uv.index]]
 name = "pytorch-rocm61"
 url = "https://download.pytorch.org/whl/rocm6.1"
@@ -201,6 +211,7 @@ torch = [
     { index = "pytorch-cu118", extra = "cu118" },
     { index = "pytorch-cu121", extra = "cu121" },
     { index = "pytorch-cu124", extra = "cu124" },
+    { index = "pytorch-cu130", extra = "cu130" },
     { index = "pytorch-rocm61", extra = "rocm" },
     { index = "pytorch-xpu", extra = "xpu" },
 ]
@@ -209,6 +220,7 @@ torchvision = [
     { index = "pytorch-cu118", extra = "cu118" },
     { index = "pytorch-cu121", extra = "cu121" },
     { index = "pytorch-cu124", extra = "cu124" },
+    { index = "pytorch-cu130", extra = "cu130" },
     { index = "pytorch-rocm61", extra = "rocm" },
     { index = "pytorch-xpu", extra = "xpu" },
 ]

src/anomalib/__init__.py

@@ -86,3 +86,28 @@ class TaskType(str, Enum):
 
     CLASSIFICATION = "classification"
     SEGMENTATION = "segmentation"
+
+
+class PrecisionType(str, Enum):
+    """Precision type defining the numerical precision used in model computations.
+
+    This enum defines the different precision types supported by anomalib models:
+
+        - ``FLOAT32``: Standard 32-bit floating point precision
+        - ``FLOAT16``: Half-precision 16-bit floating point for faster computation
+
+    Example:
+        >>> from anomalib import PrecisionType
+        >>> precision_type = PrecisionType.FLOAT16
+        >>> print(precision_type)
+        'float16'
+
+    Note:
+        The precision type affects:
+            - Memory usage during model training and inference
+            - Computational speed, especially on compatible hardware
+            - Numerical stability of model computations
+    """
+
+    FLOAT32 = "float32"
+    FLOAT16 = "float16"

src/anomalib/models/image/patchcore/lightning_model.py

@@ -56,7 +56,7 @@
 from torch import nn
 from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, Resize
 
-from anomalib import LearningType
+from anomalib import LearningType, PrecisionType
 from anomalib.data import Batch
 from anomalib.metrics import Evaluator
 from anomalib.models.components import AnomalibModule, MemoryBankMixin
@@ -93,6 +93,9 @@ class Patchcore(MemoryBankMixin, AnomalibModule):
             subsample embeddings. Defaults to ``0.1``.
         num_neighbors (int, optional): Number of nearest neighbors to use.
             Defaults to ``9``.
+        precision (str, optional): Precision type for model computations.
+            Supported values are defined in :class:`PrecisionType`.
+            Defaults to ``PrecisionType.FLOAT32``.
         pre_processor (PreProcessor | bool, optional): Pre-processor instance or
             bool flag. Defaults to ``True``.
         post_processor (PostProcessor | bool, optional): Post-processor instance or
@@ -102,6 +105,7 @@ class Patchcore(MemoryBankMixin, AnomalibModule):
         visualizer (Visualizer | bool, optional): Visualizer instance or bool flag.
             Defaults to ``True``.
 
+
     Example:
         >>> from anomalib.data import MVTecAD
         >>> from anomalib.models import Patchcore
@@ -112,7 +116,8 @@ class Patchcore(MemoryBankMixin, AnomalibModule):
         >>> model = Patchcore(
         ...     backbone="wide_resnet50_2",
         ...     layers=["layer2", "layer3"],
-        ...     coreset_sampling_ratio=0.1
+        ...     coreset_sampling_ratio=0.1,
+        ...     precision="float32",
         ... )
 
         >>> # Train using the Engine
@@ -140,6 +145,7 @@ def __init__(
         pre_trained: bool = True,
         coreset_sampling_ratio: float = 0.1,
         num_neighbors: int = 9,
+        precision: str = PrecisionType.FLOAT32,
         pre_processor: nn.Module | bool = True,
         post_processor: nn.Module | bool = True,
         evaluator: Evaluator | bool = True,
@@ -160,6 +166,15 @@ def __init__(
         )
         self.coreset_sampling_ratio = coreset_sampling_ratio
 
+        if precision == PrecisionType.FLOAT16:
+            self.model = self.model.half()
+        elif precision == PrecisionType.FLOAT32:
+            self.model = self.model.float()
+        else:
+            msg = f"""Unsupported precision type: {precision}.
+            Supported types are: {PrecisionType.FLOAT16}, {PrecisionType.FLOAT32}."""
+            raise ValueError(msg)
+
     @classmethod
     def configure_pre_processor(
         cls,