Feature Extractors

Feature extractors.

class anomalib.models.components.feature_extractors.BackboneParams(class_path: str | type[nn.Module], init_args: dict = <factory>)[source]

Bases: object

Used for serializing the backbone.

class_path: str | type[nn.Module]
init_args: dict
class anomalib.models.components.feature_extractors.FeatureExtractor(*args, **kwargs)[source]

Bases: TimmFeatureExtractor

Compatibility wrapper for the old FeatureExtractor class.

See anomalib.models.components.feature_extractors.timm.TimmFeatureExtractor for more details.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

training: bool
class anomalib.models.components.feature_extractors.TimmFeatureExtractor(backbone: str, layers: list[str], pre_trained: bool = True, requires_grad: bool = False)[source]

Bases: 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.

  • pre_trained (bool) – Whether to use a pre-trained backbone. Defaults to True.

  • requires_grad (bool) – Whether to require gradients for the backbone. Defaults to False. Models like stfpm use the feature extractor model as a trainable network. In such cases gradient computation is required.

Example

>>> import torch
>>> from anomalib.models.components.feature_extractors import TimmFeatureExtractor

>>> model = TimmFeatureExtractor(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])]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → dict[str, Tensor][source]

Forward-pass input tensor into the CNN.

Parameters:

inputs (Tensor) – Input tensor

Returns:

Feature map extracted from the CNN

training: bool
class anomalib.models.components.feature_extractors.TorchFXFeatureExtractor(backbone: str | BackboneParams | dict | nn.Module, return_nodes: list[str], weights: str | WeightsEnum | None = None, requires_grad: bool = False)[source]

Bases: Module

Extract features from a CNN.

Parameters:

  • backbone (str | BackboneParams | dict | nn.Module) – The backbone to which the feature extraction hooks are attached. If the name is provided, the model is loaded from torchvision. Otherwise, the model class can be provided and it will try to load the weights from the provided weights file. Last, an instance of nn.Module can also be passed directly.

  • return_nodes (Iterable[str]) – List of layer names of the backbone to which the hooks are attached. You can find the names of these nodes by using get_graph_node_names function.

  • weights (str | WeightsEnum | None) – Weights enum to use for the model. Torchvision models require WeightsEnum. These enums are defined in torchvision.models.<model>. You can pass the weights path for custom models.

  • requires_grad (bool) – Models like stfpm use the feature extractor for training. In such cases we should set requires_grad to True. Default is False.

Example

With torchvision models:

>>> import torch
>>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
>>> from torchvision.models.efficientnet import EfficientNet_B5_Weights
>>> feature_extractor = TorchFXFeatureExtractor(
        backbone="efficientnet_b5",
        return_nodes=["features.6.8"],
        weights=EfficientNet_B5_Weights.DEFAULT
    )
>>> input = torch.rand((32, 3, 256, 256))
>>> features = feature_extractor(input)
>>> [layer for layer in features.keys()]
    ["features.6.8"]
>>> [feature.shape for feature in features.values()]
    [torch.Size([32, 304, 8, 8])]

With custom models:

>>> import torch
>>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
>>> feature_extractor = TorchFXFeatureExtractor(
        "path.to.CustomModel", ["linear_relu_stack.3"], weights="path/to/weights.pth"
    )
>>> input = torch.randn(1, 1, 28, 28)
>>> features = feature_extractor(input)
>>> [layer for layer in features.keys()]
    ["linear_relu_stack.3"]

with model instances:

>>> import torch
>>> from anomalib.models.components.feature_extractors import TorchFXFeatureExtractor
>>> from timm import create_model
>>> model = create_model("resnet18", pretrained=True)
>>> feature_extractor = TorchFXFeatureExtractor(model, ["layer1"])
>>> input = torch.rand((32, 3, 256, 256))
>>> features = feature_extractor(input)
>>> [layer for layer in features.keys()]
    ["layer1"]
>>> [feature.shape for feature in features.values()]
    [torch.Size([32, 64, 64, 64])]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → dict[str, Tensor][source]

Extract features from the input.

initialize_feature_extractor(backbone: BackboneParams | nn.Module, return_nodes: list[str], weights: str | WeightsEnum | None = None, requires_grad: bool = False) → GraphModule[source]

Extract features from a CNN.

Parameters:

  • backbone (BackboneParams | nn.Module) – The backbone to which the feature extraction hooks are attached. If the name is provided for BackboneParams, the model is loaded from torchvision. Otherwise, the model class can be provided and it will try to load the weights from the provided weights file. Last, an instance of the model can be provided as well, which will be used as-is.

  • return_nodes (Iterable[str]) – List of layer names of the backbone to which the hooks are attached. You can find the names of these nodes by using get_graph_node_names function.

  • weights (str | WeightsEnum | None) – Weights enum to use for the model. Torchvision models require WeightsEnum. These enums are defined in torchvision.models.<model>. You can pass the weights path for custom models.

  • requires_grad (bool) – Models like stfpm use the feature extractor for training. In such cases we should set requires_grad to True. Default is False.

Returns:

Feature Extractor based on TorchFX.

training: bool
anomalib.models.components.feature_extractors.dryrun_find_featuremap_dims(feature_extractor: FeatureExtractor | GraphModule, input_size: tuple[int, int], layers: list[str]) → dict[str, dict[str, int | tuple[int, int]]][source]

Dry run an empty image of input_size size to get the featuremap tensors’ dimensions (num_features, resolution).

Returns:

maping of layer -> dimensions dict

Each dimension dict has two keys: num_features (int) and `resolution`(tuple[int, int]).

Return type:

tuple[int, int]