add precision modifier

Author: alfieroddanintel

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

@@ -53,7 +53,7 @@
 from torch import nn
 from torchvision.transforms.v2 import Compose, InterpolationMode, 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
@@ -91,6 +91,9 @@ class AnomalyDINO(MemoryBankMixin, AnomalibModule):
             to reduce the size of the memory bank. Defaults to ``False``.
         sampling ratio(float, optional): If coreset subsampling, by what ratio
             should we subsample. Defaults to ``0.1``
+        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 to enable default preprocessing. Defaults to ``True``.
         post_processor (PostProcessor | bool, optional): Post-processor instance or
@@ -152,6 +155,7 @@ def __init__(
         masking: bool = False,
         coreset_subsampling: bool = False,
         sampling_ratio: float = 0.1,
+        precision: str = PrecisionType.FLOAT32,
         pre_processor: nn.Module | bool = True,
         post_processor: nn.Module | bool = True,
         evaluator: Evaluator | bool = True,
@@ -171,6 +175,15 @@ def __init__(
             sampling_ratio=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,

src/anomalib/models/image/anomaly_dino/torch_model.py

@@ -220,6 +220,10 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor | InferenceBatch:
                     * ``pred_score``: Image-level anomaly score ``(B, 1)``
                     * ``anomaly_map``: Pixel-level anomaly heatmap ``(B, 1, H, W)``
         """
+        # set precicion
+        input_tensor = input_tensor.type(self.memory_bank.dtype)
+
+        # work out sizing
         b, _, w, h = input_tensor.shape
         cropped_width = w - w % self.feature_encoder.patch_size
         cropped_height = h - h % self.feature_encoder.patch_size