How to use custom trained model as backbone?

[lathashree01]

Describe the bug

Hi,
Thanks for anomalib library.

I am able to do Patchcore based anomaly detection, but the performance seems to be stagnant at some point.
I am playing with parameters:
number of neighbours, coreset ratio, model architecture, making dataset diverse.
I am also using pretrained models.

However, I thought if I can use a model trained specifically on my dataset. How can I load a resnet34 model pretrined on my dataset for feature extraction?

Dataset

Folder

Model

PatchCore

Steps to reproduce the behavior

How to use custom pretrained model

OS information

OS information:

• OS: [e.g. Ubuntu 20.04]
• Python version: [e.g. 3.10.0]
• Anomalib version: [e.g. 0.3.6]
• PyTorch version: [e.g. 1.9.0]
• CUDA/cuDNN version: [e.g. 11.1]
• GPU models and configuration: [e.g. 2x GeForce RTX 3090]
• Any other relevant information: [e.g. I'm using a custom dataset]

Expected behavior

Able to load model pretrained on custom dataset, for ex: using pretrained resnet34 model

Screenshots

NA

Pip/GitHub

GitHub

What version/branch did you use?

1.1.0

Configuration YAML

using engine.train function required parameters

Logs

NA

Code of Conduct

• I agree to follow this project's Code of Conduct

[samet-akcay]

Hi @lathashree01, you could refer to this documentation section for custom bacbones
https://anomalib.readthedocs.io/en/latest/markdown/guides/how_to/models/feature_extractors.html

Let us know if there is anything that is not clear

[lathashree01]

Thanks for the reply. I am able to change the backbones as shown in the link.

I want to load my custom pretrained model from local path.

When I looked in code, I found a place where pretrained_cfg is used which helps load from local path here

anomalib/src/anomalib/models/components/feature_extractors/timm.py

Line 57 in c60b5ce

if "__AT__" in backbone:

.

I am assuming that when I pass 'resnet34__AT__/path/to/checkpoint.pt' its should work. But I dont think its working for me.

When I give
backbone = "resnet34__AT__" +str(os.path.join(RESULTS_FOLDER_PATH, 'Patchcore', 'mydata', 'v12', 'weights', 'lightning' ,'model.ckpt'))
I am getting error as

    pretrained_cfg = timm.models.registry.get_pretrained_cfg(backbone)
AttributeError: module 'timm.models' has no attribute 'registry'

Can you please help resolve or point me in the right direction? Thanks

[lathashree01]

The problem is with the timm version, I had as per the anomalib requirements but it was still breaking. Please check the toml file for code compatibility. Thanks.

Downgraded to 0.6.7, the registry error is gone but still not able to load the custom model from local path.

[lathashree01]

I think the bug is,
as stated in timm package we can load the model like this - huggingface/pytorch-image-models#1826 (comment)

whereas this is not possible with current anomalib implementation.
Please let me know if my understanding is right

[blaz-r]

The timm feature extractor within the anomalib is not really designed to work with custom models I believe. TorchFX can do that, but that would require modification of PatchCore model to use TorchFX feature extractor.

[samet-akcay]

Yeah, we didn't really considers custom models for feature extraction. We thought timm or torchfx feature extractors are usually of interest.

What I would do for now is to directly customize patchcore and modify this bit however you like

anomalib/src/anomalib/models/image/patchcore/torch_model.py

Lines 48 to 52 in debdae7

	self.feature_extractor = TimmFeatureExtractor(
	backbone=self.backbone,
	pre_trained=pre_trained,
	layers=self.layers,
	).eval()

Meanwhile, we could think about how to address this

[lathashree01]

Understood. Thanks both for your responses.

Though timm models were mostly enough for prototyping and developing initial models. I wanted to explore if my custom pre-trained model might yield better performance. My problem domain has completely different setup from imagenet dataset, hence the thought.

[blaz-r]

You could try the torchfx extractor with your model https://anomalib.readthedocs.io/en/latest/markdown/guides/reference/models/components/feature_extractors.html#anomalib.models.components.feature_extractors.TorchFXFeatureExtractor

Instead of the timm feature extractor on the lines that Samet linked above. I didn't work that much with this specific anomalib version but in theory you can load the model from the path and pass it to the feature extractor directly. Just make sure you get the right layer names (docs also cover this).

[watertianyi]

@blaz-r @samet-akcay
I used convnextv2_huge.fcmae_ft_in22k_in1k_512 to replace the pre-trained network wide_resnet50_2 in patchcore, and used layers=['stages.1', 'stages.2'] for classification. The indicator effect is far worse than wide_resnet50_2. Isn't convnextv2 better than wide_resnet50_2? Why Would it be so bad?

[alexriedel1]

Did you pre-train for defect image classification? Did you also train WideResNet on the same task and compared both after?
Did you try different layers when using your pre-trained convnext than stage 1 and stage 2?

There is some research on how different backbones work in patchcore:
https://arxiv.org/pdf/2307.10792

[watertianyi]

Yes, I trained convnext using defective data, but I didn't train WideResNet using defective data because I used convnext and used

 # self.model = Patchcore(
        #     backbone="convnextv2_pico.fcmae",
        #     pre_trained=False,
        #     coreset_sampling_ratio=0.1,
        #     layers=['stages.1', 'stages.2', 'stages.3'],
        # )

[alexriedel1]

Ok, so then you can't compare the two if they are trained differently. I suspect you didn't train your convnext properly and thats why it probably performs bad. Additionally ConvNext 2 pico performs worse on ImageNet compared to Wide ResNet 50

[alexriedel1]

self.model = Patchcore(
# backbone="convnextv2_pico.fcmae",
# pre_trained=False,
# coreset_sampling_ratio=0.1,
# layers=['stages.1', 'stages.2', 'stages.3'],
# )

If this is all the modifications you used, then you are training patchcore on an randomly initiliazed convnext model. Where are you loading your custom pre-trained model?

[watertianyi]

`self.feature_extractor = timm.create_model(
            backbone,
            pretrained=pre_trained,
            pretrained_cfg=pretrained_cfg,
            features_only=True,
            exportable=True,
            out_indices=self.idx,
            in_chans =in_channels,
        )
checkpoint_path = ""
checkpoint = torch.load(checkpoint_path)
checkpoint_model = checkpoint['model']
f_checkpoint = checkpoint_filter_fn(checkpoint_model,model=self.feature_extractor)
self.feature_extractor.load_state_dict(f_checkpoint,strict=False)`

`def checkpoint_filter_fn(state_dict, model=None):
    """ Remap FB checkpoints -> timm """
    if 'head.norm.weight' in state_dict or 'norm_pre.weight' in state_dict:
        return state_dict  # non-FB checkpoint
    if 'model' in state_dict:
        state_dict = state_dict['model']
        print("model:")

    out_dict = {}
    if 'visual.trunk.stem.0.weight' in state_dict:
        print('visual.trunk.stem.0.weight')
        out_dict = {k.replace('visual.trunk.', ''): v for k, v in state_dict.items() if k.startswith('visual.trunk.')}
        if 'visual.head.proj.weight' in state_dict:
            out_dict['head.fc.weight'] = state_dict['visual.head.proj.weight']
            out_dict['head.fc.bias'] = torch.zeros(state_dict['visual.head.proj.weight'].shape[0])
        elif 'visual.head.mlp.fc1.weight' in state_dict:
            out_dict['head.pre_logits.fc.weight'] = state_dict['visual.head.mlp.fc1.weight']
            out_dict['head.pre_logits.fc.bias'] = state_dict['visual.head.mlp.fc1.bias']
            out_dict['head.fc.weight'] = state_dict['visual.head.mlp.fc2.weight']
            out_dict['head.fc.bias'] = torch.zeros(state_dict['visual.head.mlp.fc2.weight'].shape[0])
        return out_dict

    import re
    for k, v in state_dict.items():
        if 'downsample_layers.0.' in k:
            print('downsample_layers.0.:',k,v.shape)
        k = k.replace('downsample_layers.0.', 'stem.')
        
        k = re.sub(r'stages.([0-9]+).([0-9]+)', r'stages.\1.blocks.\2', k)
        k = re.sub(r'downsample_layers.([0-9]+).([0-9]+)', r'stages.\1.downsample.\2', k)
        k = k.replace('dwconv', 'conv_dw')
        k = k.replace('pwconv', 'mlp.fc')
        if 'grn' in k:
            k = k.replace('grn.beta', 'mlp.grn.bias')
            k = k.replace('grn.gamma', 'mlp.grn.weight')
            v = v.reshape(v.shape[-1])
        k = k.replace('head.', 'head.fc.')
        if k.startswith('norm.'):
            k = k.replace('norm', 'head.norm')
        if v.ndim == 2 and 'head' not in k:
            print("v.dim:",v.ndim,k,model.state_dict()[k].shape)
            model_shape = model.state_dict()[k].shape
            v = v.reshape(model_shape)
        out_dict[k] = v
    
    return out_dict`