7047 simplify resnet pretrained

monai/networks/nets/resnet.py

@@ -14,6 +14,10 @@
 from collections.abc import Callable
 from functools import partial
 from typing import Any
+from pathlib import Path
+import logging
+import re
+from huggingface_hub import hf_hub_download
 
 import torch
 import torch.nn as nn
@@ -23,6 +27,9 @@
 from monai.utils import ensure_tuple_rep
 from monai.utils.module import look_up_option
 
+MEDICALNET_HUGGINGFACE_REPO_BASENAME = "TencentMedicalNet/MedicalNet-Resnet"
+MEDICALNET_HUGGINGFACE_FILES_BASENAME = "resnet_"
+
 __all__ = [
     "ResNet",
     "ResNetBlock",
@@ -36,6 +43,9 @@
     "resnet200",
 ]
 
+logger = logging.getLogger(__name__)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_inplanes():
     return [64, 128, 256, 512]
@@ -329,21 +339,40 @@ def _resnet(
     block: type[ResNetBlock | ResNetBottleneck],
     layers: list[int],
     block_inplanes: list[int],
-    pretrained: bool,
+    pretrained: bool | str,
     progress: bool,
     **kwargs: Any,
 ) -> ResNet:
     model: ResNet = ResNet(block, layers, block_inplanes, **kwargs)
     if pretrained:
-        # Author of paper zipped the state_dict on googledrive,
-        # so would need to download, unzip and read (2.8gb file for a ~150mb state dict).
-        # Would like to load dict from url but need somewhere to save the state dicts.
-        raise NotImplementedError(
-            "Currently not implemented. You need to manually download weights provided by the paper's author"
-            " and load then to the model with `state_dict`. See https://github.com/Tencent/MedicalNet"
-            "Please ensure you pass the appropriate `shortcut_type` and `bias_downsample` args. as specified"
-            "here: https://github.com/Tencent/MedicalNet/tree/18c8bb6cd564eb1b964bffef1f4c2283f1ae6e7b#update20190730"
-        )
+        if isinstance(pretrained, str):
+            if Path(pretrained).exists():
+                logger.info(f"Loading weights from {pretrained}...")
+                checkpoint = torch.load(pretrained, map_location=device)
+            else:
+                ### Throw error
+                raise FileNotFoundError("The pretrained checkpoint file is not found")
+        else:
+            resnet_depth = int(re.search(r"resnet(\d+)", arch).group(1))
+            # Download the MedicalNet pretrained model
+            logger.info(f"Loading MedicalNet pretrained model from https://huggingface.co/{MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}")
+            pretrained_path = hf_hub_download(
+                repo_id=f"{MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}",
+                filename=f"{MEDICALNET_HUGGINGFACE_FILES_BASENAME}{resnet_depth}.pth")
+
+            checkpoint = torch.load(pretrained_path, map_location=device)
+
+        if "state_dict" in checkpoint:
+            model_state_dict = checkpoint["state_dict"]
+            model_state_dict = {key.replace("module.", ""): value for key, value in model_state_dict.items()}
+        else:
+            ### Throw error
+            raise KeyError(
+                "The checkpoint should contain the pretrained model state dict with the following key: 'state_dict'"
+            )
+
+        model.load_state_dict(model_state_dict, strict=True)
+
     return model
 
 

tests/test_resnet.py

@@ -12,6 +12,8 @@
 from __future__ import annotations
 
 import unittest
+import copy
+import os
 from typing import TYPE_CHECKING
 
 import torch
@@ -30,6 +32,8 @@
 else:
     torchvision, has_torchvision = optional_import("torchvision")
 
+# from torchvision.models import ResNet50_Weights, resnet50
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 TEST_CASE_1 = [  # 3D, batch 3, 2 input channel
@@ -159,9 +163,11 @@
 ]
 
 TEST_CASES = []
+PRETRAINED_TEST_CASES = []
 for case in [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_2_A, TEST_CASE_3_A]:
     for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200]:
         TEST_CASES.append([model, *case])
+        PRETRAINED_TEST_CASES.append([model, *case])
 for case in [TEST_CASE_5, TEST_CASE_5_A, TEST_CASE_6, TEST_CASE_7]:
     TEST_CASES.append([ResNet, *case])
 
@@ -181,6 +187,27 @@ def test_resnet_shape(self, model, input_param, input_shape, expected_shape):
             else:
                 self.assertTrue(result.shape in expected_shape)
 
+    @parameterized.expand(PRETRAINED_TEST_CASES)
+    def test_resnet_pretrained(self, model, input_param, input_shape, expected_shape):
+        net = model(**input_param).to(device)
+        tmp_ckpt_filename = "monai_unittest_tmp_ckpt.pth"
+        # Save ckpt
+        torch.save({
+            "state_dict": net.state_dict()
+        },
+        tmp_ckpt_filename)
+
+        cp_input_param = copy.copy(input_param)
+        cp_input_param["pretrained"] = tmp_ckpt_filename
+        pretrained_net = model(**cp_input_param)
+        assert str(net.state_dict()) == str(pretrained_net.state_dict())
+
+        with self.assertRaises(NotImplementedError):
+            cp_input_param["pretrained"] = True
+            model(**cp_input_param)
+
+        os.remove(tmp_ckpt_filename)
+
     @parameterized.expand(TEST_SCRIPT_CASES)
     def test_script(self, model, input_param, input_shape, expected_shape):
         net = model(**input_param)