Anomalib Adding A Model

Adding a New Model

Models live under src/anomalib/models/image/<model_name>/ (or .../video/<model_name>/ for video models). Every model is a LightningModule subclass of AnomalibModule (src/anomalib/models/components/base/anomalib_module.py) that wraps a plain torch.nn.Module.

Reference implementation

Read src/anomalib/models/image/padim/ first — it is the smallest complete example:

• torch_model.py — PadimModel(nn.Module): pure PyTorch forward pass. In training mode it returns intermediate embeddings; in eval mode it returns an InferenceBatch(pred_score=..., anomaly_map=...).
• anomaly_map.py — AnomalyMapGenerator(nn.Module): post-processing (score/map computation) kept out of torch_model.py for clarity. Not all models need a separate file for this.
• lightning_model.py — Padim(MemoryBankMixin, AnomalibModule): the Lightning-facing wrapper. This is the class users construct (Padim()), pass to Engine, and reference from CLI/config as anomalib.models.Padim.
• __init__.py — exports the Lightning class only: from .lightning_model import Padim.
• README.md — usage + benchmark notes (see model-doc-sync/model-sample-image-export skills for docs work).

Only mix in MemoryBankMixin (src/anomalib/models/components/base/memory_bank_module.py) if the model accumulates a memory bank / feature bank across training (as PaDiM and PatchCore do). Most models just subclass AnomalibModule directly.

Required members on your AnomalibModule subclass

import torch
from anomalib import LearningType
from anomalib.data import Batch
from anomalib.metrics import Evaluator
from anomalib.models.components import AnomalibModule
from anomalib.post_processing import PostProcessor
from anomalib.pre_processing import PreProcessor
from anomalib.visualization import Visualizer

class MyModel(AnomalibModule):
    def __init__(self, some_param: int = 1, pre_processor: PreProcessor | bool = True,
                 post_processor: PostProcessor | bool = True,
                 evaluator: Evaluator | bool = True,
                 visualizer: Visualizer | bool = True) -> None:
        super().__init__(pre_processor=pre_processor, post_processor=post_processor,
                          evaluator=evaluator, visualizer=visualizer)
        self.model = MyTorchModel(some_param=some_param)

    @property
    def trainer_arguments(self) -> dict:
        """Default Trainer overrides for this model, e.g. {"max_epochs": 1} for training-free models."""
        return {"max_epochs": 1, "num_sanity_val_steps": 0}

    @property
    def learning_type(self) -> LearningType:
        return LearningType.ONE_CLASS

    def training_step(self, batch: Batch, *args, **kwargs) -> torch.Tensor:
        # Training-free models (e.g. Padim) still return a dummy loss for Lightning.
        _ = self.model(batch.image)
        return torch.tensor(0.0, requires_grad=True, device=self.device)

    def validation_step(self, batch: Batch, *args, **kwargs) -> Batch:
        predictions = self.model(batch.image)
        return batch.update(**predictions._asdict())

    def configure_optimizers(self) -> None:
        # Return None for training-free / statistical models (Padim does this).
        return None

AnomalibModule provides working defaults you can override only when the model needs something different:

• configure_pre_processor(image_size=None) — resize + ImageNet normalization.
• configure_post_processor() — thresholding/normalization for ONE_CLASS models.
• configure_evaluator() — image/pixel AUROC and F1 metrics.
• configure_visualizer() — default ImageVisualizer.

pre_processor / post_processor / evaluator / visualizer constructor args each accept an instance, True (use the configured default), or False (disable).

Registration — how the model becomes discoverable

1. Add the export in src/anomalib/models/image/__init__.py (or video/__init__.py):

   from .my_model import MyModel
   

   and add "MyModel" to __all__ and to the Available Models docstring list at the top of the file.

2. Also add the import in src/anomalib/models/__init__.py, which explicitly imports all image (and video) models and defines the top-level __all__. Without this, anomalib.models.MyModel won't resolve.

3. That's the only registration needed. list_models() and get_model() (src/anomalib/models/__init__.py) discover models by walking AnomalibModule.__subclasses__() and matching on cls.__name__ (case-insensitive) — there is no separate name-string registry to update.

4. get_model() also accepts a dict/DictConfig/Namespace with class_path + init_args, restricted to modules listed in ALLOWED_MODULES (same file) — you don't need to touch that set for models already under anomalib.models.

5. Once exported, the model is usable as anomalib.models.MyModel, from get_model("MyModel"), and from the CLI: anomalib train --model MyModel --data anomalib.data.MVTecAD.

Tests

Add tests/unit/models/image/my_model/test_my_model.py (or .../video/my_model/...). At minimum:

• get_model("MyModel") (and the equivalent PascalCase/snake_case variants) returns a MyModel instance.
• model.trainer_arguments is a dict and model.learning_type is the expected LearningType.
• A synthetic batch through training_step/validation_step produces the expected shapes: image-level pred_score has shape (batch_size,) and pixel-level anomaly_map has shape (batch_size, 1, H, W) without crashing.

See tests/unit/models/test_model_utils.py for the get_model() instantiation pattern used across the suite.

Gotchas

• Keep torch_model.py importable and testable without Lightning — the nn.Module.forward contract (train-mode returns raw tensors/embeddings, eval-mode returns InferenceBatch) is what the base class and visualizers expect. Don't put Lightning-specific logic there.
• Training side effects (checkpointing, timing, compression, visualization) belong in src/anomalib/callbacks/-style Lightning callbacks or existing hooks — not inline in the model's training_step.
• Public constructor arguments must stay explicit and typed (no untyped **kwargs passthrough) so jsonargparse can expose them on the CLI/config surface.
• If the model becomes part of the public API, update the model's README.md and, if one exists, the matching page under docs/source/markdown/guides/reference/models/.

Reviewer / self-check before opening a PR

• Model exported from src/anomalib/models/image/__init__.py (or video/__init__.py) and __all__ updated.
• trainer_arguments, learning_type, training_step, validation_step, configure_optimizers implemented.
• get_model("MyModel") and anomalib.models.MyModel both resolve.
• Unit tests added under tests/unit/models/.
• README.md added in the model folder.
• pre-commit run --all-files and pytest tests/unit/models/ -k my_model pass.