STFPM#

Student-Teacher Feature Pyramid Matching for anomaly detection.

This module implements the STFPM model for anomaly detection as described in Wang et al. (2021).

The model consists of: - A pre-trained teacher network that extracts multi-scale features - A student network that learns to match the teacher’s feature representations - Feature pyramid matching between student and teacher features - Anomaly detection based on feature discrepancy

Example

>>> from anomalib.models.image import Stfpm
>>> from anomalib.engine import Engine
>>> from anomalib.data import MVTec
>>> datamodule = MVTec()
>>> model = Stfpm(
...     backbone="resnet18",
...     layers=["layer1", "layer2", "layer3"]
... )
>>> engine = Engine(model=model, datamodule=datamodule)
>>> engine.fit()  
>>> predictions = engine.predict()  

See also

  • Stfpm: Lightning implementation of the model

  • STFPMModel: PyTorch implementation of the model architecture

  • STFPMLoss: Loss function for training

class anomalib.models.image.stfpm.lightning_model.Stfpm(backbone='resnet18', layers=('layer1', 'layer2', 'layer3'), pre_processor=True, post_processor=True, evaluator=True, visualizer=True)#

Bases: AnomalibModule

PL Lightning Module for the STFPM algorithm.

The Student-Teacher Feature Pyramid Matching (STFPM) model consists of a pre-trained teacher network and a student network that learns to match the teacher’s feature representations. The model detects anomalies by comparing feature discrepancies between the teacher and student networks.

Parameters:
  • backbone (str) – Name of the backbone CNN network used for both teacher and student. Defaults to "resnet18".

  • layers (list[str]) – Names of layers from which to extract features. Defaults to ["layer1", "layer2", "layer3"].

  • pre_processor (PreProcessor | bool, optional) – Pre-processor to transform input data before passing to model. If True, uses default. Defaults to True.

  • post_processor (PostProcessor | bool, optional) – Post-processor to generate predictions from model outputs. If True, uses default. Defaults to True.

  • evaluator (Evaluator | bool, optional) – Evaluator to compute metrics. If True, uses default. Defaults to True.

  • visualizer (Visualizer | bool, optional) – Visualizer to display results. If True, uses default. Defaults to True.

Example

>>> from anomalib.models.image import Stfpm
>>> from anomalib.data import MVTec
>>> from anomalib.engine import Engine
>>> datamodule = MVTec()
>>> model = Stfpm(
...     backbone="resnet18",
...     layers=["layer1", "layer2", "layer3"]
... )
>>> engine = Engine(model=model, datamodule=datamodule)
>>> engine.fit()  
>>> predictions = engine.predict()  

See also

configure_optimizers()#

Configure optimizers for training.

Returns:

SGD optimizer with the following parameters:
  • Learning rate: 0.4

  • Momentum: 0.9

  • Weight decay: 0.001

Return type:

torch.optim.Optimizer

property learning_type: LearningType#

Get the learning type of the model.

Returns:

The model uses one-class learning.

Return type:

LearningType

property trainer_arguments: dict[str, Any]#

Get required trainer arguments for the model.

Returns:

Dictionary of trainer arguments:
  • gradient_clip_val: Set to 0 to disable gradient clipping

  • num_sanity_val_steps: Set to 0 to skip validation sanity checks

Return type:

dict[str, Any]

training_step(batch, *args, **kwargs)#

Perform a training step of STFPM.

For each batch, teacher and student features are extracted from the CNN.

Parameters:
  • batch (Batch) – Input batch containing images and labels.

  • args – Additional arguments (unused).

  • kwargs – Additional keyword arguments (unused).

Returns:

Dictionary containing the loss value.

Return type:

STEP_OUTPUT

validation_step(batch, *args, **kwargs)#

Perform a validation step of STFPM.

Similar to training, extracts student/teacher features from CNN and computes anomaly maps.

Parameters:
  • batch (Batch) – Input batch containing images and labels.

  • args – Additional arguments (unused).

  • kwargs – Additional keyword arguments (unused).

Returns:

Dictionary containing images, anomaly maps, labels and

masks for evaluation.

Return type:

STEP_OUTPUT

PyTorch model implementation for Student-Teacher Feature Pyramid Matching.

This module implements the core PyTorch model architecture for the STFPM anomaly detection method as described in Wang et al. (2021).

The model consists of: - A pre-trained teacher network that extracts multi-scale features - A student network that learns to match the teacher’s feature representations - Feature pyramid matching between student and teacher features - Anomaly detection based on feature discrepancy

Example

>>> from anomalib.models.image.stfpm.torch_model import STFPMModel
>>> model = STFPMModel(
...     backbone="resnet18",
...     layers=["layer1", "layer2", "layer3"]
... )
>>> features = model(torch.randn(1, 3, 256, 256))

See also

  • STFPMModel: Main PyTorch model implementation

  • STFPMLoss: Loss function for training

  • AnomalyMapGenerator: Anomaly map generation from features

class anomalib.models.image.stfpm.torch_model.STFPMModel(layers, backbone='resnet18')#

Bases: Module

PyTorch implementation of the STFPM model.

The Student-Teacher Feature Pyramid Matching model consists of a pre-trained teacher network and a student network that learns to match the teacher’s feature representations. The model detects anomalies by comparing feature discrepancies between the teacher and student networks.

Parameters:
  • layers (Sequence[str]) – Names of layers from which to extract features. For example ["layer1", "layer2", "layer3"].

  • backbone (str, optional) – Name of the backbone CNN architecture used for both teacher and student networks. Supported backbones can be found in timm library. Defaults to "resnet18".

Example

>>> import torch
>>> from anomalib.models.image.stfpm.torch_model import STFPMModel
>>> model = STFPMModel(
...     backbone="resnet18",
...     layers=["layer1", "layer2", "layer3"]
... )
>>> input_tensor = torch.randn(1, 3, 256, 256)
>>> features = model(input_tensor)

Note

The teacher model is initialized with pre-trained weights and frozen during training, while the student model is trained from scratch.

tiler#

Optional tiler for processing large images in patches.

Type:

Tiler | None

teacher_model#

Pre-trained teacher network for feature extraction.

Type:

TimmFeatureExtractor

student_model#

Student network that learns to match teacher features.

Type:

TimmFeatureExtractor

anomaly_map_generator#

Module to generate anomaly maps from features.

Type:

AnomalyMapGenerator

forward(images)#

Forward pass through teacher and student networks.

The forward pass behavior differs between training and evaluation: - Training: Returns features from both teacher and student networks - Evaluation: Returns anomaly maps generated from feature differences

Parameters:

images (torch.Tensor) – Batch of input images with shape (N, C, H, W).

Returns:

tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:

Features from teacher and student networks respectively. Each dict maps layer names to feature tensors.

Evaluation mode:
InferenceBatch:

Batch containing anomaly maps and prediction scores.

Return type:

Training mode

Example

>>> import torch
>>> from anomalib.models.image.stfpm.torch_model import STFPMModel
>>> model = STFPMModel(layers=["layer1", "layer2", "layer3"])
>>> input_tensor = torch.randn(1, 3, 256, 256)
>>> # Training mode
>>> model.train()
>>> teacher_feats, student_feats = model(input_tensor)
>>> # Evaluation mode
>>> model.eval()
>>> predictions = model(input_tensor)

Loss function for Student-Teacher Feature Pyramid Matching model.

This module implements the loss function used to train the STFPM model for anomaly detection as described in Wang et al. (2021).

The loss function: 1. Takes feature maps from teacher and student networks as input 2. Normalizes the features using L2 normalization 3. Computes MSE loss between normalized features 4. Scales the loss by spatial dimensions of feature maps

Example

>>> from anomalib.models.components import TimmFeatureExtractor
>>> from anomalib.models.image.stfpm.loss import STFPMLoss
>>> from torchvision.models import resnet18
>>> layers = ["layer1", "layer2", "layer3"]
>>> teacher_model = TimmFeatureExtractor(
...     model=resnet18(pretrained=True),
...     layers=layers
... )
>>> student_model = TimmFeatureExtractor(
...     model=resnet18(pretrained=False),
...     layers=layers
... )
>>> criterion = STFPMLoss()
>>> features = torch.randn(4, 3, 256, 256)
>>> teacher_features = teacher_model(features)
>>> student_features = student_model(features)
>>> loss = criterion(student_features, teacher_features)

See also

  • STFPMLoss: Main loss class implementation

  • Stfpm: Lightning implementation of the full model

class anomalib.models.image.stfpm.loss.STFPMLoss#

Bases: Module

Loss function for Student-Teacher Feature Pyramid Matching model.

This class implements the feature pyramid loss function proposed in the STFPM paper. The loss measures the discrepancy between feature representations from a pre-trained teacher network and a student network that learns to match them.

The loss computation involves: 1. Normalizing teacher and student features using L2 normalization 2. Computing MSE loss between normalized features 3. Scaling the loss by spatial dimensions of feature maps 4. Summing losses across all feature layers

Example

>>> from anomalib.models.components import TimmFeatureExtractor
>>> from anomalib.models.image.stfpm.loss import STFPMLoss
>>> from torchvision.models import resnet18
>>> layers = ["layer1", "layer2", "layer3"]
>>> teacher_model = TimmFeatureExtractor(
...     model=resnet18(pretrained=True),
...     layers=layers
... )
>>> student_model = TimmFeatureExtractor(
...     model=resnet18(pretrained=False),
...     layers=layers
... )
>>> criterion = STFPMLoss()
>>> features = torch.randn(4, 3, 256, 256)
>>> teacher_features = teacher_model(features)
>>> student_features = student_model(features)
>>> loss = criterion(student_features, teacher_features)
>>> loss
tensor(51.2015, grad_fn=<SumBackward0>)

See also

  • Stfpm: Lightning implementation of the full model

  • STFPMModel: PyTorch implementation of the model architecture

compute_layer_loss(teacher_feats, student_feats)#

Compute loss between teacher and student features for a single layer.

This implements the loss computation based on Equation (1) in Section 3.2 of the paper. The loss is computed as: 1. L2 normalize teacher and student features 2. Compute MSE loss between normalized features 3. Scale loss by spatial dimensions (height * width)

Parameters:
  • teacher_feats (torch.Tensor) – Features from teacher network with shape (B, C, H, W)

  • student_feats (torch.Tensor) – Features from student network with shape (B, C, H, W)

Returns:

Scalar loss value for the layer

Return type:

torch.Tensor

forward(teacher_features, student_features)#

Compute total loss across all feature layers.

The total loss is computed as the sum of individual layer losses. Each layer loss measures the discrepancy between teacher and student features at that layer.

Parameters:
  • teacher_features (dict[str, torch.Tensor]) – Dictionary mapping layer names to teacher feature tensors

  • student_features (dict[str, torch.Tensor]) – Dictionary mapping layer names to student feature tensors

Returns:

Total loss summed across all layers

Return type:

torch.Tensor

Anomaly map computation for Student-Teacher Feature Pyramid Matching model.

This module implements functionality to generate anomaly heatmaps by comparing features between a pre-trained teacher network and a student network that learns to match the teacher’s representations.

The anomaly maps are generated by: 1. Computing cosine similarity between teacher and student features 2. Converting similarity scores to anomaly scores via L2 norm 3. Upscaling anomaly scores to original image size 4. Combining multiple layer scores via element-wise multiplication

Example

>>> from anomalib.models.image.stfpm.anomaly_map import AnomalyMapGenerator
>>> generator = AnomalyMapGenerator()
>>> teacher_features = {"layer1": torch.randn(1, 64, 32, 32)}
>>> student_features = {"layer1": torch.randn(1, 64, 32, 32)}
>>> anomaly_map = generator.compute_anomaly_map(
...     teacher_features,
...     student_features,
...     image_size=(256, 256)
... )

See also

  • AnomalyMapGenerator: Main class for generating anomaly maps

  • compute_layer_map(): Function to compute per-layer anomaly scores

class anomalib.models.image.stfpm.anomaly_map.AnomalyMapGenerator#

Bases: Module

Generate anomaly heatmaps by comparing teacher and student features.

This class implements functionality to generate anomaly maps by comparing feature representations between a pre-trained teacher network and a student network. The comparison is done via cosine similarity and L2 distance.

The anomaly map generation process involves: 1. Computing cosine similarity between teacher-student feature pairs 2. Converting similarity scores to anomaly scores using L2 norm 3. Upscaling the scores to original image size 4. Combining multiple layer scores via element-wise multiplication

Example

>>> from anomalib.models.image.stfpm.anomaly_map import AnomalyMapGenerator
>>> generator = AnomalyMapGenerator()
>>> teacher_features = {"layer1": torch.randn(1, 64, 32, 32)}
>>> student_features = {"layer1": torch.randn(1, 64, 32, 32)}
>>> anomaly_map = generator.compute_anomaly_map(
...     teacher_features,
...     student_features,
...     image_size=(256, 256)
... )

See also

compute_anomaly_map(teacher_features, student_features, image_size)#

Compute overall anomaly map by combining multiple layer maps.

The final anomaly map is generated by: 1. Computing per-layer anomaly maps via compute_layer_map() 2. Combining layer maps through element-wise multiplication

Parameters:
  • teacher_features (dict[str, torch.Tensor]) – Dictionary mapping layer names to teacher feature tensors

  • student_features (dict[str, torch.Tensor]) – Dictionary mapping layer names to student feature tensors

  • image_size (tuple[int, int] | torch.Size) – Target size for the anomaly map in format (H, W)

Returns:

Final anomaly map with shape (B, 1, H, W) where

B is batch size and (H, W) matches image_size

Return type:

torch.Tensor

static compute_layer_map(teacher_features, student_features, image_size)#

Compute anomaly map for a single feature layer.

The layer map is computed by: 1. Normalizing teacher and student features 2. Computing L2 distance between normalized features 3. Upscaling the distance map to original image size

Parameters:
  • teacher_features (torch.Tensor) – Features from teacher network with shape (B, C, H, W)

  • student_features (torch.Tensor) – Features from student network with matching shape

  • image_size (tuple[int, int] | torch.Size) – Target size for upscaling in format (H, W)

Returns:

Anomaly scores for the layer, upscaled to

image_size

Return type:

torch.Tensor

forward(**kwargs)#

Generate anomaly map from teacher and student features.

Expects the following keys in kwargs: - teacher_features: Dictionary of teacher network features - student_features: Dictionary of student network features - image_size: Target size for the anomaly map

Parameters:

kwargs (dict[str, torch.Tensor]) – Keyword arguments containing required inputs

Example

>>> generator = AnomalyMapGenerator()
>>> anomaly_map = generator(
...     teacher_features=teacher_features,
...     student_features=student_features,
...     image_size=(256, 256)
... )
Raises:

ValueError – If required keys are missing from kwargs

Returns:

Anomaly map with shape (B, 1, H, W)

Return type:

torch.Tensor