Fix pretrained weight cached dir for timm (#3909)

* Fix pretrained_weight for timm

* Fix unit-test

Author: harimkang

src/otx/algo/classification/backbones/timm.py

@@ -9,14 +9,13 @@
 """
 from __future__ import annotations
 
-from pathlib import Path
 from typing import Literal
 
 import timm
 import torch
 from torch import nn
 
-from otx.algo.utils.mmengine_utils import load_checkpoint_to_model, load_from_http
+from otx.algo.utils.mmengine_utils import load_from_http
 
 PRETRAINED_ROOT = "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/"
 pretrained_urls = {
@@ -28,7 +27,7 @@
     "mobilenetv3_large_21k": "mobilenetv3_large_100_miil_in21k",
     "mobilenetv3_large_1k": "mobilenetv3_large_100_miil",
     "tresnet": "tresnet_m",
-    "efficientnetv2_s_21k": "tf_efficientnetv2_s_in21k",
+    "efficientnetv2_s_21k": "tf_efficientnetv2_s.in21k",
     "efficientnetv2_s_1k": "tf_efficientnetv2_s_in21ft1k",
     "efficientnetv2_m_21k": "tf_efficientnetv2_m_in21k",
     "efficientnetv2_m_1k": "tf_efficientnetv2_m_in21ft1k",
@@ -59,12 +58,19 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.backbone = backbone
-        self.pretrained = pretrained
+        self.pretrained: bool | dict = pretrained
         self.is_mobilenet = backbone.startswith("mobilenet")
+        if pretrained and self.backbone in pretrained_urls:
+            # This pretrained weight is saved into ~/.cache/torch/hub/checkpoints
+            # Otherwise, it is stored in ~/.cache/huggingface/hub. (timm defaults)
+            self.pretrained = load_from_http(filename=pretrained_urls[self.backbone])
+
+        self.model = timm.create_model(
+            TIMM_MODEL_NAME_DICT[self.backbone],
+            pretrained=self.pretrained,
+            num_classes=1000,
+        )
 
-        self.model = timm.create_model(TIMM_MODEL_NAME_DICT[self.backbone], pretrained=pretrained, num_classes=1000)
-        if self.pretrained:
-            print(f"init weight - {pretrained_urls[self.backbone]}")
         self.model.classifier = None  # Detach classifier. Only use 'backbone' part in otx.
         self.num_head_features = self.model.num_features
         self.num_features = self.model.conv_head.in_channels if self.is_mobilenet else self.model.num_features
@@ -97,15 +103,3 @@ def get_config_optim(self, lrs: list[float] | float) -> list[dict[str, float]]:
                 param_dict["lr"] = lrs
 
         return parameters
-
-    def init_weights(self, pretrained: str | bool | None = None) -> None:
-        """Initialize weights."""
-        checkpoint = None
-        if isinstance(pretrained, str) and Path(pretrained).exists():
-            checkpoint = torch.load(pretrained, None)
-            print(f"init weight - {pretrained}")
-        elif pretrained is not None:
-            checkpoint = load_from_http(pretrained_urls[self.key])
-            print(f"init weight - {pretrained_urls[self.key]}")
-        if checkpoint is not None:
-            load_checkpoint_to_model(self, checkpoint)

tests/unit/algo/classification/backbones/test_timm.py

@@ -8,7 +8,6 @@
 class TestOTXEfficientNetV2:
     def test_forward(self):
         model = TimmBackbone(backbone="efficientnetv2_s_21k")
-        model.init_weights()
         assert model(torch.randn(1, 3, 244, 244))[0].shape == torch.Size([1, 1280, 8, 8])
 
     def test_get_config_optim(self):