Feature extractor

Feature extractor#

This guide demonstrates how different backbones can be used as feature extractors for anomaly detection models. Here we show how to use API and CLI to use different backbones as feature extractors.

Backbones can be defined by their name as available in Timm or directly passed to the feature extractor as torch nn.Module.

See also

For specifics of implementation refer to implementation classes Timm Feature Extractor

Available backbones and layers#

Available Timm models are listed on Timm GitHub page.

In most cases, we want to use a pretrained backbone, so can get a list of all such models using the following code:

import timm
# list all pretrained models in timm
for model_name in timm.list_models(pretrained=True):
    print(model_name)

Once we have a model selected we can obtain available layer names using the following code:

import timm
model = timm.create_model("resnet18", features_only=True)
# Print module names
print(model.feature_info.module_name())
>>>['act1', 'layer1', 'layer2', 'layer3', 'layer4']

model = timm.create_model("mobilenetv3_large_100", features_only=True)
print(model.feature_info.module_name())
>>>['blocks.0.0', 'blocks.1.1', 'blocks.2.2', 'blocks.4.1', 'blocks.6.0']

We can then use selected model name and layer names with either API or using config file.

Warning

Some models might not support every backbone.

Backbone and layer selection#

When using API, we need to specify backbone and layers when instantiating the model with a non-default timm backbone.

 1# Import the required modules
 2from anomalib.data import MVTecAD
 3from anomalib.models import Padim
 4from anomalib.engine import Engine
 5
 6# Initialize the datamodule, model, and engine
 7datamodule = MVTecAD(num_workers=0)
 8# Specify backbone and layers
 9model = Padim(backbone="resnet18", layers=["layer1", "layer2"])
10engine = Engine(image_metrics=["AUROC"], pixel_metrics=["AUROC"])
11
12# Train the model
13engine.fit(datamodule=datamodule, model=model)

In the following example config, we can see that we need to specify two parameters: the backbone and layers list.

 1model:
 2  class_path: anomalib.models.Padim
 3  init_args:
 4    layers:
 5      - blocks.1.1
 6      - blocks.2.2
 7    input_size: null
 8    backbone: mobilenetv3_large_100
 9    pre_trained: true
10    n_features: 50

Then we can train using:

anomalib train --config <path/to/config>

Custom backbone#

To use a custom backbone model, e.g. locally saved models or models with modified weights, the model instance is passed as backbone argument. Features will be extracted from the specified layers. Make sure that the backbone model actually possesses the layers from which the feature will be extracted.

 1# Import the required modules
 2import torch
 3from torchvision import models
 4
 5from anomalib.data import MVTec
 6from anomalib.models import Padim
 7from anomalib.engine import Engine
 8
 9# Initialize the datamodule, model, and engine
10datamodule = MVTec(num_workers=0)
11
12# Specify custom model
13weights = torch.hub.load_state_dict_from_url("https://huggingface.co/mzweilin/robust-imagenet-models/resolve/main/wide_resnet50_2_l2_eps5.pth")
14custom_backbone = models.wide_resnet50_2()
15custom_backbone.load_state_dict(weights)
16
17# Specify backbone and layers
18model = Padim(backbone=custom_backbone, layers=["layer1", "layer3"])
19engine = Engine(image_metrics=["AUROC"], pixel_metrics=["AUROC"])
20
21# Train the model
22engine.fit(datamodule=datamodule, model=model)
Contents