fix(backbone): update model loading to use CPU mapping

Author: CuriousDolphin

focoos/nn/backbone/convnextv2.py

@@ -158,7 +158,7 @@ def __init__(self, config: ConvNeXtV2Config):
             cur += depths[i]
 
         if config.use_pretrained and backbone_url:
-            state = torch.hub.load_state_dict_from_url(backbone_url)
+            state = torch.hub.load_state_dict_from_url(backbone_url, map_location="cpu")
             self.load_state_dict(state)
             logger.info(f"Load ConvNeXtV2{config.model_size} state_dict")
 

focoos/nn/backbone/csp_darknet.py

@@ -369,11 +369,11 @@ def __init__(
         }
         if config.use_pretrained:
             if config.backbone_url:
-                state = torch.hub.load_state_dict_from_url(config.backbone_url)
+                state = torch.hub.load_state_dict_from_url(config.backbone_url, map_location="cpu")
                 self.load_state_dict(state)
                 logger.info(f"Loaded pretrained weights from {config.backbone_url}")
             else:
-                state = torch.hub.load_state_dict_from_url(CONFIGS[config.size]["url"])
+                state = torch.hub.load_state_dict_from_url(CONFIGS[config.size]["url"], map_location="cpu")
                 self.load_state_dict(state)
                 logger.info(f"Loaded pretrained weights from {CONFIGS[config.size]['url']}")
 

focoos/nn/backbone/mobilenet_v2.py

@@ -205,7 +205,7 @@ def __init__(
             self.layers.append(layer_name)
 
         if config.use_pretrained and config.backbone_url:
-            state = torch.hub.load_state_dict_from_url(config.backbone_url)
+            state = torch.hub.load_state_dict_from_url(config.backbone_url, map_location="cpu")
             self.load_state_dict(state)
             logger.info("Load MobileNetV2 state_dict")
 

focoos/nn/backbone/resnet.py

@@ -227,7 +227,7 @@ def __init__(
             self._freeze_norm(self)
 
         if use_pretrained:
-            state = torch.hub.load_state_dict_from_url(backbone_url)
+            state = torch.hub.load_state_dict_from_url(backbone_url, map_location="cpu")
             self.load_state_dict(state)
             logger.info(f"Load ResNet{depth} state_dict")
 

focoos/nn/backbone/stdc.py

@@ -248,7 +248,7 @@ def __init__(self, config: STDCConfig):
             self.out_ids = 1, 5, 10, 13
 
         if config.use_pretrained and config.backbone_url:
-            state = torch.hub.load_state_dict_from_url(config.backbone_url)
+            state = torch.hub.load_state_dict_from_url(config.backbone_url, map_location="cpu")
             self.load_state_dict(state)
             logger.info("Load STDC state_dict")
 

focoos/nn/backbone/swin.py

@@ -679,7 +679,7 @@ def __init__(
             self.add_module(layer_name, layer)
 
         if config.use_pretrained and backbone_url:
-            state = torch.hub.load_state_dict_from_url(backbone_url)
+            state = torch.hub.load_state_dict_from_url(backbone_url, map_location="cpu")
             self.load_state_dict(state, strict=False)
             logger.info(f"Loaded pretrained weights from {backbone_url}")
 

tests/test_backbone.py

@@ -11,16 +11,49 @@
     {"model_type": "stdc", "use_pretrained": True, "size": size} for size in ["nano", "small", "large"]
 ]
 stdc_configs_base = [
-    {"model_type": "stdc", "use_pretrained": True, "base": 64, "layers": [2, 2, 2], "block_num": 4, "block_type": "cat"}
+    {
+        "model_type": "stdc",
+        "use_pretrained": True,
+        "base": 64,
+        "layers": [2, 2, 2],
+        "block_num": 4,
+        "block_type": "cat",
+    },
+    {
+        "model_type": "stdc",
+        "use_pretrained": False,
+        "base": 64,
+        "layers": [4, 5, 3],
+        "block_num": 4,
+        "block_type": "cat",
+    },
 ]
 BACKBONE_CONFIGS = {
-    "resnet": [{"model_type": "resnet", "use_pretrained": False, "depth": 18}],
+    "resnet": [
+        {"model_type": "resnet", "use_pretrained": True, "depth": 18},
+        {"model_type": "resnet", "use_pretrained": True, "depth": 34},
+        {"model_type": "resnet", "use_pretrained": False, "depth": 50},
+        {"model_type": "resnet", "use_pretrained": False, "depth": 101},
+    ],
     "stdc": stdc_configs_size + stdc_configs_base,
-    "swin": [{"model_type": "swin", "use_pretrained": False}],
-    "mobilenet_v2": [{"model_type": "mobilenet_v2", "use_pretrained": False}],
-    "convnextv2": [{"model_type": "convnextv2", "use_pretrained": False}],
+    "swin": [
+        {"model_type": "swin", "use_pretrained": False},
+        {
+            "model_type": "swin",
+            "use_pretrained": True,
+        },
+    ],
+    "mobilenet_v2": [
+        {"model_type": "mobilenet_v2", "use_pretrained": False},
+        {"model_type": "mobilenet_v2", "use_pretrained": True},
+    ],
+    "convnextv2": [
+        {"model_type": "convnextv2", "use_pretrained": False},
+        {"model_type": "convnextv2", "use_pretrained": True},
+    ],
     "csp_darknet": [
-        {"model_type": "csp_darknet", "use_pretrained": False, "size": size} for size in ["small", "medium", "large"]
+        {"model_type": "csp_darknet", "use_pretrained": True if size == "small" else False, "size": size}
+        for size in ["small", "medium", "large"]
     ],
 }