Merge pull request #46 from DeepKnowledge1/develop

feat(inference): add backend-specific warm-up support

- Implemented warmup() in each backend:
  • TorchBackend: runs dummy or provided tensor through model with AMP/cuDNN
  • TorchScriptBackend: runs dummy or provided tensor on-device a few times
  • OnnxBackend: generates dummy numpy input matching model metadata and executes
  • OpenVinoBackend: creates infer request and runs with dummy or provided input
  • TensorRTBackend: added stub (warns and skips, since not implemented)

- Added ModelWrapper.warmup() to delegate to backend if supported
- Updated detect.py to run warm-up with first dataloader batch before inference loop

Warm-up stabilizes first-batch latency, preloads kernels, and reduces variance,
especially on CUDA backends (Torch/TorchScript/ONNX).

Author: DeepKnowledge1

anodet/inference/model/backends/onnx_backend.py

@@ -70,3 +70,22 @@ def predict(self, batch: Batch) -> ScoresMaps:
     def close(self) -> None:
         """Release the ONNX session."""
         self.session = None
+
+
+    def warmup(self, batch, runs: int = 2) -> None:
+        """
+        Warm up ONNX Runtime by running the session a few times.
+        Args:
+            batch: input batch to use for warm-up.
+            runs: Number of warm-up runs to perform.
+        """
+        if isinstance(batch, np.ndarray):
+            input_arr = batch
+        else:
+            input_arr = batch.detach().cpu().numpy()
+
+        feeds = {self.input_names[0]: input_arr}
+        for _ in range(max(1, runs)):
+            _ = self.session.run(self.output_names, feeds)
+
+        logger.info("OnnxBackend warm-up completed (runs=%d, shape=%s).", runs, tuple(input_arr.shape))

anodet/inference/model/backends/openvino_backend.py

@@ -91,3 +91,23 @@ def close(self) -> None:
         self.compiled_model = None
         self.model = None
         self.core = None
+
+    def warmup(self, batch=None, runs: int = 2) -> None:
+        """
+        Warm up OpenVINO by creating an infer request and calling infer repeatedly.
+        Args:
+            batch: input batch to use for warm-up.
+            runs: Number of warm-up runs to perform.
+        """
+
+        if isinstance(batch, np.ndarray):
+            input_arr = batch
+        else:
+            input_arr = batch.detach().cpu().numpy()
+
+        infer_request = self.compiled_model.create_infer_request()
+
+        for _ in range(max(1, runs)):
+            _ = infer_request.infer({self.input_layer.any_name: input_arr})
+
+        logger.info("OpenVinoBackend warm-up completed (runs=%d, shape=%s).", runs, tuple(input_arr.shape))

anodet/inference/model/backends/tensorrt_backend.py

@@ -24,3 +24,7 @@ def predict(self, batch: Batch) -> ScoresMaps:
 
     def close(self) -> None:
         pass
+
+    def warmup(self, batch=None, runs: int = 2) -> None:
+        logger.warning("TensorRT backend is not implemented; warm-up skipped.")
+        raise NotImplementedError("TensorRT warm-up not implemented yet.")

anodet/inference/model/backends/torch_backend.py

@@ -13,6 +13,7 @@
 from .base import Batch, ScoresMaps, InferenceBackend
 from anodet.utils import get_logger
 
+
 logger = get_logger(__name__)
 
 
@@ -77,3 +78,34 @@ def predict(self, batch: Batch) -> ScoresMaps:
     def close(self) -> None:
         """Release PyTorch model."""
         self.model = None
+
+    def warmup(self, batch, runs: int = 2) -> None:
+        """
+        Warm up PyTorch/TorchScript-on-PyTorch backend.
+        Uses AMP on CUDA if enabled, and executes the model a few times.
+        Args:
+            batch: input batch to use for warm-up.
+            runs: Number of warm-up runs to perform.
+
+        """
+        import torch
+        from contextlib import nullcontext
+
+        if not isinstance(batch, torch.Tensor):
+            batch = torch.as_tensor(batch, dtype=torch.float32, device=self.device)
+        else:
+            batch = batch.to(self.device, non_blocking=True)
+
+        # Avoid autograd + use AMP if configured
+        autocast_ctx = (
+            torch.autocast(device_type=self.device.type, dtype=torch.float16)
+            if self.use_amp and self.device.type == "cuda"
+            else nullcontext()
+        )
+
+        with torch.inference_mode(), autocast_ctx:
+            for _ in range(max(1, runs)):
+                # Prefer model.predict to match your runtime path
+                _ = self.model.predict(batch)
+
+        logger.info("TorchBackend warm-up completed (runs=%d, shape=%s).", runs, tuple(batch.shape))

anodet/inference/model/backends/torchscript_backend.py

@@ -71,3 +71,23 @@ def close(self) -> None:
         self.model = None
         if self.device.type == "cuda":
             torch.cuda.empty_cache()
+
+    def warmup(self, batch=None, runs: int = 2) -> None:
+        """
+        Warm up TorchScript backend by forwarding a few times on-device.
+        Args:
+            batch: input batch to use for warm-up.
+            runs: Number of warm-up runs to perform.
+
+        """
+
+        if isinstance(batch, torch.Tensor):
+            batch = batch.to(self.device, non_blocking=True)
+        else:
+            batch = torch.as_tensor(batch, dtype=torch.float32, device=self.device)
+
+        with torch.no_grad():
+            for _ in range(max(1, runs)):
+                _ = self.model(batch)
+
+        logger.info("TorchScriptBackend warm-up completed (runs=%d, shape=%s).", runs, tuple(batch.shape))

anodet/inference/model/wrapper.py

@@ -95,3 +95,9 @@ def close(self) -> None:
         """Release resources associated with the backend."""
         logger.info("Closing ModelWrapper and releasing resources")
         self.backend.close()
+
+    def warmup(self, batch=None, runs: int = 2) -> None:
+        if hasattr(self.backend, "warmup"):
+            return self.backend.warmup(batch=batch, runs=runs)
+        else:
+            logger.info(f"{self.backend.__class__.__name__} does not support warm-up. Skipping.")

detect.py

@@ -14,7 +14,10 @@
 import numpy as np
 import torch
 from torch.utils.data import DataLoader
+import matplotlib
+matplotlib.use("Agg")  # non-interactive, faster PNG writing
 import matplotlib.pyplot as plt
+
 import argparse
 import time
 from datetime import datetime
@@ -46,7 +49,7 @@ def parse_args():
     parser.add_argument(
         "--model",
         type=str,
-        default="padim_model.onnx",
+        default="padim_model.pt",
         help="Model file (.pt for PyTorch, .onnx for ONNX, .engine for TensorRT)",
     )
     parser.add_argument(
@@ -85,7 +88,7 @@ def parse_args():
     )
     parser.add_argument(
         "--save_visualizations",
-        action="store_false",
+        action="store_true",
         help="Save visualization images to disk.",
     )
     parser.add_argument(
@@ -108,12 +111,6 @@ def parse_args():
     )
 
 
-    parser.add_argument(
-        "--show_first_batch_only",
-        action="store_true",
-        default=True,
-        help="Show visualization only for the first batch.",
-    )
     parser.add_argument(
         "--viz_alpha", type=float, default=0.5, help="Alpha value for heatmap overlay."
     )
@@ -231,6 +228,22 @@ def main(args):
 
     logger.info(f"Processing {len(test_dataset)} images using AnomaVision {model_type.value.upper()}")
 
+    # ---- Warm-up
+    try:
+        first = next(iter(test_dataloader))  # (batch, images, _, _)
+        first_batch = first[0].to(device_str)
+        model.warmup(batch=first_batch, runs=2)
+        logger.info("AnomaVision warm-up done with first batch %s.", tuple(first_batch.shape))
+    except StopIteration:
+        logger.warning("Dataset empty; skipping warm-up.")
+    except Exception as e:
+        logger.warning(f"Warm-up skipped due to error: {e}")
+
+    # ---- End warm-up
+
+
+
+
     # AnomaVision batch processing pipeline
     batch_count = 0
     try:
@@ -306,36 +319,37 @@ def main(args):
                         if args.save_visualizations:
                             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
 
-                        # Display AnomaVision results for first batch only (if requested)
-                        if batch_idx == 0 and (not args.show_first_batch_only or batch_idx == 0):
+                        # Display AnomaVision results
+                        for img_id in range(len(images)):
                             try:
                                 fig, axs = plt.subplots(1, 4, figsize=(16, 8))
-                                fig.suptitle(f"AnomaVision Detection Results - Batch {batch_idx + 1}", fontsize=14)
+                                fig.suptitle(f"AnomaVision Detection Results - Batch {img_id + 1}", fontsize=14)
 
-                                axs[0].imshow(images[0])
+                                axs[0].imshow(images[img_id])
                                 axs[0].set_title("Original Image")
                                 axs[0].axis("off")
 
-                                axs[1].imshow(boundary_images[0])
+                                axs[1].imshow(boundary_images[img_id])
                                 axs[1].set_title("AnomaVision Boundary Detection")
                                 axs[1].axis("off")
 
-                                axs[2].imshow(heatmap_images[0])
+                                axs[2].imshow(heatmap_images[img_id])
                                 axs[2].set_title("AnomaVision Anomaly Heatmap")
                                 axs[2].axis("off")
 
-                                axs[3].imshow(highlighted_images[0])
+                                axs[3].imshow(highlighted_images[img_id])
                                 axs[3].set_title("AnomaVision Highlighted Anomalies")
                                 axs[3].axis("off")
 
-                                plt.tight_layout()
+                                # plt.tight_layout()
 
                                 if args.save_visualizations:
                                     combined_filepath = os.path.join(RESULTS_PATH, f"anomavision_batch_{batch_idx}_{timestamp}.png")
-                                    plt.savefig(combined_filepath, dpi=300, bbox_inches="tight")
+                                    plt.savefig(combined_filepath, dpi=100, bbox_inches="tight")
                                     logger.info(f"AnomaVision visualization saved: {combined_filepath}")
 
-                                plt.show()
+                                # plt.show()
+                                plt.close(fig)
 
                             except Exception as e:
                                 logger.warning(f"Failed to display AnomaVision visualization for batch {batch_idx}: {e}")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "AnomaVision"
-version = "2.0.30"
+version = "2.0.37"
 description = "Deep learnIng Anomaly Detection EnvironMent [AnomaVision] is a deep learning library that aims to collect state-of-the-art anomaly detection algorithms for benchmarking on both public and private datasets. PaDimOpti provides several ready-to-use implementations of anomaly detection algorithms described in the recent literature, as well as a set of tools that facilitate the development and implementation of custom models. The library has a strong focus on image-based anomaly detection, where the goal of the algorithm is to identify anomalous images, or anomalous pixel regions within images in a dataset. PaDimOpti is constantly being updated with new algorithms and training/inference extensions, so stay tuned!!"
 authors = ["Deep Knowledge <Deepp.Knowledge@gmail.com>"]
 readme = "README.md"