Improve backbone weights loading mechanism

CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Improve backbone weights loading mechanism by @mzweilin in https://github.com/openvinotoolkit/anomalib/pull/2454
+
 ### New Contributors
 
 ## [v1.2.0]

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

@@ -53,12 +53,10 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        # Extract backbone-name and weight-URI from the backbone string.
+        # Extract the backbone name and the weight file location from the backbone string.
         if "__AT__" in backbone:
-            backbone, uri = backbone.split("__AT__")
-            pretrained_cfg = timm.models.registry.get_pretrained_cfg(backbone)
-            # Override pretrained_cfg["url"] to use different pretrained weights.
-            pretrained_cfg["url"] = uri
+            backbone, location = backbone.split("__AT__")
+            pretrained_cfg = {"url": location} if location.startswith(("http://", "https://")) else {"file": location}
         else:
             pretrained_cfg = None
 

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

@@ -82,7 +82,9 @@ def __init__(
             pre_trained=pre_trained,
         ).eval()
         self.n_features_original = sum(self.feature_extractor.out_dims)
-        self.n_features = n_features or _N_FEATURES_DEFAULTS.get(self.backbone)
+        # In case the backbone has the weight file information.
+        backbone_name = self.backbone.split("__AT__")[0]
+        self.n_features = n_features or _N_FEATURES_DEFAULTS.get(backbone_name)
         if self.n_features is None:
             msg = (
                 f"n_features must be specified for backbone {self.backbone}. "

tests/unit/models/components/base/test_feature_extractor.py

@@ -0,0 +1,37 @@
+"""Unit tests for TimmFeatureExtractor."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import tempfile
+from pathlib import Path
+
+import torch
+from timm.models import create_model
+
+from anomalib.models.components.feature_extractors.timm import TimmFeatureExtractor
+
+
+def test_backbone_weight_file() -> None:
+    """Test the backbone weight file loading mechanism."""
+    # Use the simplest model.
+    backbone = "resnet18"
+    _, state_dict_fpath = tempfile.mkstemp()
+    # Only examine conv1 before layers in the feature extractor.
+    layers = []
+
+    # Get random model weights without downloading.
+    model = create_model(backbone, pretrained=False)
+    state_dict = model.state_dict()
+
+    # Set the conv1 weight to zero, and save state_dict to a temp file.
+    state_dict["conv1.weight"].zero_()
+    torch.save(state_dict, state_dict_fpath)
+
+    # Load weights from the temp file.
+    backbone_with_path = f"{backbone}__AT__{state_dict_fpath}"
+
+    fe_restored = TimmFeatureExtractor(backbone_with_path, layers, pre_trained=True)
+    Path(state_dict_fpath).unlink()
+    # The weights should be zero if the file loading mechanism works.
+    assert torch.all(fe_restored.feature_extractor.conv1.weight == 0)