UniNet

UniNet#

Lightning model for UniNet.

class anomalib.models.image.uninet.lightning_model.UniNet(student_backbone='wide_resnet50_2', teacher_backbone='wide_resnet50_2', temperature=0.1, pre_processor=True, post_processor=True, evaluator=True, visualizer=True)#

Bases: AnomalibModule

UniNet model for anomaly detection.

Parameters:

  • student_backbone (str) – The backbone model to use for the student network. Defaults to “wide_resnet50_2”.

  • teacher_backbone (str) – The backbone model to use for the teacher network. Defaults to “wide_resnet50_2”.

  • temperature (float) – Temperature parameter used for contrastive loss. Controls the temperature of the student and teacher similarity computation. Defaults to 0.1.

  • pre_processor (PreProcessor | bool) – Preprocessor instance or bool flag. Defaults to True.

  • post_processor (PostProcessor | bool) – Postprocessor instance or bool flag. Defaults to True.

  • evaluator (Evaluator | bool) – Evaluator instance or bool flag. Defaults to True.

  • visualizer (Visualizer | bool) – Visualizer instance or bool flag. Defaults to True.

configure_optimizers()#

Configure optimizers for training.

Returns:

Optimizers for student and target teacher.

Return type:

tuple[Optimizer, Optimizer]

property learning_type: LearningType#

The model uses one-class learning.

property trainer_arguments: dict[str, Any]#

Does not require any trainer arguments.

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

Perform a training step of UniNet.

Return type:

Union[Tensor, Mapping[str, Any], None]

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

Perform a validation step of UniNet.

Return type:

Union[Tensor, Mapping[str, Any], None]

PyTorch model for UniNet.

See also

anomalib.models.image.uninet.lightning_model.UniNet:

UniNet Lightning model.

class anomalib.models.image.uninet.torch_model.Teachers(teacher_backbone)#

Bases: Module

Teachers module for UniNet.

Parameters:

  • source_teacher (nn.Module) – Source teacher model.

  • target_teacher (nn.Module | None) – Target teacher model.

forward(images)#

Forward pass of the teachers.

Parameters:

images (Tensor) – Input images.

Returns:

Source features or source and target features.

Return type:

tuple[Tensor, list[Tensor]]

class anomalib.models.image.uninet.torch_model.UniNetModel(student_backbone, teacher_backbone, loss)#

Bases: Module

UniNet PyTorch model.

It consists of teachers, student, and bottleneck modules.

Parameters:

  • student_backbone (str) – Student backbone model.

  • teacher_backbone (str) – Teacher backbone model.

  • loss (Module) – Loss function.

forward(images, masks=None, labels=None)#

Forward pass of the UniNet model.

Parameters:
Returns:

Loss or InferenceBatch.

Return type:

Tensor | InferenceBatch

Loss functions for UniNet.

class anomalib.models.image.uninet.components.loss.UniNetLoss(lambda_weight=0.7, temperature=2.0)#

Bases: Module

Loss function for UniNet.

Parameters:

  • lambda_weight (float) – Hyperparameter for balancing the loss. Defaults to 0.7.

  • temperature (float) – Temperature for contrastive learning. Defaults to 2.0.

forward(student_features, teacher_features, margin=1, mask=None, stop_gradient=False)#

Compute the loss.

Parameters:

  • student_features (list[Tensor]) – Student features.

  • teacher_features (list[Tensor]) – Teacher features.

  • margin (int) – Hyperparameter for controlling the boundary.

  • mask (Tensor | None) – Mask for the prediction. Mask is of shape Bx1xHxW

  • stop_gradient (bool) – Whether to stop the gradient into teacher features.

Return type:

Tensor

Utilities for computing anomaly maps.

anomalib.models.image.uninet.components.anomaly_map.weighted_decision_mechanism(batch_size, output_list, alpha, beta, output_size=(256, 256))#

Compute anomaly maps using weighted decision mechanism.

Parameters:

  • batch_size (int) – Batch size.

  • output_list (list[Tensor]) – List of output tensors, each with shape [batch_size, H, W].

  • alpha (float) – Alpha parameter. Used for controlling the upper limit

  • beta (float) – Beta parameter. Used for controlling the lower limit

  • output_size (tuple[int, int]) – Output size.

Returns:

Anomaly score and anomaly map.

Return type:

tuple[Tensor, Tensor]

Attention Bottleneck for UniNet.

class anomalib.models.image.uninet.components.attention_bottleneck.AttentionBottleneck(inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, norm_layer=None, attention=True, halve=1)#

Bases: Module

Attention Bottleneck block for UniNet with dual-branch processing.

This module implements a specialized bottleneck block that can operate in two modes: - Standard bottleneck (when halve=1) - Dual-branch attention mechanism (when halve=2)

The dual-branch mode splits input channels and processes them through different kernel sizes (3x3 and 7x7) to capture features at different receptive field scales, then fuses them before the final expansion layer.

Architecture:

  • Input: inplanes channels

  • Intermediate: width channels (compressed for efficiency)

  • Output: planes * expansion channels (expanded representation)

Parameters:

  • inplanes (int) – Number of input channels.

  • planes (int) – Base number of channels for intermediate processing. Final output will have planes * expansion channels.

  • stride (int) – Stride for convolution layers. Defaults to 1.

  • downsample (Module | None) – Module for downsampling the residual connection when dimensions don’t match. Defaults to None.

  • groups (int) – Number of blocked connections from input to output channels. Defaults to 1.

  • base_width (int) – Base width for calculating intermediate channel width. Defaults to 64.

  • norm_layer (Callable[..., Module] | None) – Normalization layer to use. If None, uses BatchNorm2d. Defaults to None.

  • attention (bool) – Whether to use attention mechanism. Defaults to True.

  • halve (int) – Controls processing mode: - 1: Standard bottleneck processing - 2: Dual-branch processing with 3x3 and 7x7 kernels Defaults to 1.

channel_expansion#

Channel expansion factor. Final output channels will be planes * channel_expansion. Set to 4 following ResNet conventions.

Type:

int

Example

>>> import torch
>>> from anomalib.models.image.uninet.components.attention_bottleneck import (
...     AttentionBottleneck
... )
>>> # Standard bottleneck
>>> block = AttentionBottleneck(256, 64, halve=1)
>>> x = torch.randn(32, 256, 28, 28)
>>> output = block(x)
>>> output.shape  # Output: 32, 256, 28, 28 (64 * 4 = 256)
torch.Size([32, 256, 28, 28])

>>> # Dual-branch attention bottleneck
>>> block = AttentionBottleneck(512, 128, halve=2)
>>> x = torch.randn(32, 512, 14, 14)
>>> output = block(x)
>>> output.shape  # Output: 32, 512, 14, 14 (128 * 4 = 512)
torch.Size([32, 512, 14, 14])

Notes

  • When halve=2, input is split into two branches processed by 3x3 and 7x7 convolutions respectively to capture multi-scale features

  • The merge_kernel method can fuse batch norm parameters into convolution weights for inference optimization

  • Residual connections are used following ResNet design principles

  • The width calculation: int(planes * (base_width / 64.0)) * groups

  • The channel_expansion is set to 4 following ResNet conventions.

See also

  • BottleneckLayer: Container for multiple AttentionBottleneck blocks

  • fuse_bn(): Function for fusing batch normalization into convolution layers

forward(x)#

Forward pass of the bottleneck.

Parameters:

x (Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

Tensor

get_same_kernel_bias()#

Get same kernel and bias of the bottleneck.

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

merge_kernel()#

Merge kernel of the bottleneck.

Return type:

None

class anomalib.models.image.uninet.components.attention_bottleneck.BottleneckLayer(block, layers, groups=1, width_per_group=64, norm_layer=None, halve=2)#

Bases: Module

Batch Normalization layer for UniNet.

Parameters:

  • block (type[AttentionBottleneck]) – Attention bottleneck.

  • layers (int) – Number of layers.

  • groups (int) – Number of groups.

  • width_per_group (int) – Width per group.

  • norm_layer (Callable[..., Module] | None) – Normalization layer.

  • halve (int) – Number of halved channels.

forward(x)#

Forward pass of the bottleneck.

Parameters:

x (Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

Tensor

anomalib.models.image.uninet.components.attention_bottleneck.conv1x1(in_planes, out_planes, stride=1)#

1x1 convolution.

Parameters:

  • in_planes (int) – Number of input planes.

  • out_planes (int) – Number of output planes.

  • stride (int) – Stride of the convolution.

Return type:

Conv2d

anomalib.models.image.uninet.components.attention_bottleneck.conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1)#

3x3 convolution with padding.

Parameters:

  • in_planes (int) – Number of input planes.

  • out_planes (int) – Number of output planes.

  • stride (int) – Stride of the convolution.

  • groups (int) – Number of groups.

  • dilation (int) – Dilation rate.

Returns:

3x3 convolution with padding.

Return type:

Conv2d

anomalib.models.image.uninet.components.attention_bottleneck.fuse_bn(conv, bn)#

Fuse convolution and batch normalization layers.

Parameters:

  • conv (Module) – Convolution layer.

  • bn (Module) – Batch normalization layer.

Returns:

Fused convolution and batch normalization layers.

Return type:

tuple[Tensor, Tensor]

Domain Related Feature Selection.

class anomalib.models.image.uninet.components.dfs.DomainRelatedFeatureSelection(num_channels=256, learnable=True)#

Bases: Module

Domain Related Feature Selection.

It is used to select the domain-related features from the source and target features.

Parameters:

  • num_channels (int) – Number of channels in the features. Defaults to 256.

  • learnable (bool) – Whether to use learnable theta. Theta controls the domain-related feature selection.

forward(source_features, target_features, conv=False, maximize=True)#

Domain related feature selection.

Parameters:

  • source_features (list[Tensor]) – Source features.

  • target_features (list[Tensor]) – Target features.

  • conv (bool) – Whether to use convolutional domain-related feature selection. Defaults to False.

  • maximize (bool) – Used for weights computation. If True, the weights are computed by subtracting the max value from the target feature. Defaults to True.

Returns:

Domain related features.

Return type:

list[Tensor]