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 modelSTFPMModel
: PyTorch implementation of the model architectureSTFPMLoss
: 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 toTrue
.post_processor (PostProcessor | bool, optional) – Post-processor to generate predictions from model outputs. If
True
, uses default. Defaults toTrue
.evaluator (Evaluator | bool, optional) – Evaluator to compute metrics. If
True
, uses default. Defaults toTrue
.visualizer (Visualizer | bool, optional) – Visualizer to display results. If
True
, uses default. Defaults toTrue
.
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
anomalib.models.image.stfpm.torch_model.STFPMModel
:PyTorch implementation of the model architecture
anomalib.models.image.stfpm.loss.STFPMLoss
:Loss function for training
- 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:
- property learning_type: LearningType#
Get the learning type of the model.
- Returns:
The model uses one-class learning.
- Return type:
LearningType
- 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 implementationSTFPMLoss
: Loss function for trainingAnomalyMapGenerator
: 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:
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.
- teacher_model#
Pre-trained teacher network for feature extraction.
- Type:
- student_model#
Student network that learns to match teacher features.
- Type:
- anomaly_map_generator#
Module to generate anomaly maps from features.
- Type:
- 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 implementationStfpm
: 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 modelSTFPMModel
: 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:
- 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:
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 mapscompute_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_layer_map()
: Function to compute per-layer anomaly scorescompute_anomaly_map()
: Function to combine layer scores
- 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)
matchesimage_size
- Final anomaly map with shape
- Return type:
- 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:
- 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: