anomalib.models.components.feature_extractors.feature_extractor

Feature Extractor.

This script extracts features from a CNN network

Module Contents

Classes

FeatureExtractor

Extract features from a CNN.

class anomalib.models.components.feature_extractors.feature_extractor.FeatureExtractor(backbone: str, layers: List[str], pre_trained: bool = True)[source]

Bases: torch.nn.Module

Extract features from a CNN.

Parameters
  • backbone (nn.Module) – The backbone to which the feature extraction hooks are attached.

  • layers (Iterable[str]) – List of layer names of the backbone to which the hooks are attached.

Example

>>> import torch
>>> from anomalib.core.model.feature_extractor import FeatureExtractor
>>> model = FeatureExtractor(model="resnet18", layers=['layer1', 'layer2', 'layer3'])
>>> input = torch.rand((32, 3, 256, 256))
>>> features = model(input)
>>> [layer for layer in features.keys()]
    ['layer1', 'layer2', 'layer3']
>>> [feature.shape for feature in features.values()]
    [torch.Size([32, 64, 64, 64]), torch.Size([32, 128, 32, 32]), torch.Size([32, 256, 16, 16])]
_map_layer_to_idx(offset: int = 3) List[int][source]

Maps set of layer names to indices of model.

Parameters

offset (int) –

Returns

Feature map extracted from the CNN

forward(input_tensor: torch.Tensor) Dict[str, torch.Tensor][source]

Forward-pass input tensor into the CNN.

Parameters

input_tensor (Tensor) – Input tensor

Returns

Feature map extracted from the CNN