anomalib.models.reverse_distillation.lightning_model

Anomaly Detection via Reverse Distillation from One-Class Embedding.

https://arxiv.org/abs/2201.10703v2

Module Contents

Classes

ReverseDistillation

PL Lightning Module for Reverse Distillation Algorithm.

ReverseDistillationLightning

PL Lightning Module for Reverse Distillation Algorithm.
class anomalib.models.reverse_distillation.lightning_model.ReverseDistillation(input_size: Tuple[int, int], backbone: str, layers: List[str], anomaly_map_mode: str, lr: float, beta1: float, beta2: float, pre_trained: bool = True)[source]

Bases: anomalib.models.components.AnomalyModule

PL Lightning Module for Reverse Distillation Algorithm.

Parameters

  • input_size (Tuple[int, int]) – Size of model input

  • backbone (str) – Backbone of CNN network

  • layers (List[str]) – Layers to extract features from the backbone CNN

  • pre_trained (bool, optional) – Boolean to check whether to use a pre_trained backbone.

configure_optimizers()[source]

Configures optimizers for decoder and bottleneck.

Note

This method is used for the existing CLI. When PL CLI is introduced, configure optimizers method will be

deprecated, and optimizers will be configured from either config.yaml file or from CLI.

Returns

Adam optimizer for each decoder

Return type

Optimizer

training_step(batch, _) → Dict[str, torch.Tensor][source]

Training Step of Reverse Distillation Model.

Features are extracted from three layers of the Encoder model. These are passed to the bottleneck layer that are passed to the decoder network. The loss is then calculated based on the cosine similarity between the encoder and decoder features.

Parameters

  • batch (Tensor) – Input batch

  • _ – Index of the batch.

Returns

Feature Map

validation_step(batch, _)[source]

Validation Step of Reverse Distillation Model.

Similar to the training step, encoder/decoder features are extracted from the CNN for each batch, and anomaly map is computed.

Parameters

  • batch (Tensor) – Input batch

  • _ – Index of the batch.

Returns

Dictionary containing images, anomaly maps, true labels and masks. These are required in validation_epoch_end for feature concatenation.

class anomalib.models.reverse_distillation.lightning_model.ReverseDistillationLightning(hparams: Union[omegaconf.DictConfig, omegaconf.ListConfig])[source]

Bases: ReverseDistillation

PL Lightning Module for Reverse Distillation Algorithm.

Parameters

hparams (Union[DictConfig, ListConfig]) – Model parameters

configure_callbacks()[source]

Configure model-specific callbacks.

Note

This method is used for the existing CLI. When PL CLI is introduced, configure callback method will be

deprecated, and callbacks will be configured from either config.yaml file or from CLI.