7047 simplify resnet pretrained

monai/networks/nets/resnet.py

@@ -14,14 +14,23 @@
 from collections.abc import Callable
 from functools import partial
 from typing import Any
+import logging
+
 
 import torch
 import torch.nn as nn
 
 from monai.networks.layers.factories import Conv, Norm, Pool
 from monai.networks.layers.utils import get_pool_layer
-from monai.utils import ensure_tuple_rep
+from monai.utils import ensure_tuple_rep, optional_import
 from monai.utils.module import look_up_option
+from monai.networks.utils import get_pretrained_resnet_medicalnet
+
+Path, _ = optional_import("pathlib", name = "Path")
+re, _ = optional_import("re")
+
+MEDICALNET_HUGGINGFACE_REPO_BASENAME = "TencentMedicalNet/MedicalNet-Resnet"
+MEDICALNET_HUGGINGFACE_FILES_BASENAME = "resnet_"
 
 __all__ = [
     "ResNet",
@@ -36,6 +45,9 @@
     "resnet200",
 ]
 
+logger = logging.getLogger(__name__)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_inplanes():
     return [64, 128, 256, 512]
@@ -329,24 +341,51 @@ 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}...")
+                model_state_dict = torch.load(pretrained, map_location=device)
+            else:
+                ### Throw error
+                raise FileNotFoundError("The pretrained checkpoint file is not found")
+        else:
+            # Also check bias downsample and shortcut.
+            if kwargs.get("spatial_dims", 3) == 3:
+                if kwargs.get("n_input_channels", 3)==1 and kwargs.get("feed_forward", True) is False:
+                    resnet_depth = int(re.search(r"resnet(\d+)", arch).group(1))
+                    # get shortcut_type and bias_downsample.
+                    def get_medicalnet_pretrained_resnet_args(resnet_depth: int) :
+                        """
+                        Return correct shortcut_type and bias_downsample for pretrained MedicalNet weights according to rensnet depth
+                        """
+                        # After testing
+                        # False: 10, 50, 101, 152, 200
+                        # Any: 18, 34
+                        bias_downsample = -1 if resnet_depth in [18, 34] else 0  # 18, 10, 34
+                        shortcut_type = "A" if resnet_depth in [18, 34] else "B"
+                        return bias_downsample, shortcut_type
+
+                    # Check model bias_downsample and shortcut_type
+                    bias_downsample, shortcut_type = get_medicalnet_pretrained_resnet_args(resnet_depth)
+                    if shortcut_type == kwargs.get("shortcut_type", "B") and (bool(bias_downsample) == kwargs.get("bias_downsample", False) if bias_downsample != -1 else True):
+                        # Download the MedicalNet pretrained model
+                        model_state_dict = get_pretrained_resnet_medicalnet(resnet_depth, device=device, datasets23=True)
+                    else:
+                        raise NotImplementedError(f"Please set shortcut_type to {shortcut_type} and bias_downsample to {bool(bias_downsample) if bias_downsample!=-1 else 'True or False'} when using pretrained MedicalNet resnet{resnet_depth}")
+                else:
+                    raise NotImplementedError("Please set n_input_channels to 1 and feed_forward to False in order to use MedicalNet pretrained weights")
+            else:
+                raise NotImplementedError("MedicalNet pretrained weights are only avalaible for 3D models")
+        model_state_dict = {key.replace("module.", ""): value for key, value in model_state_dict.items()}
+        model.load_state_dict(model_state_dict, strict=True)
     return model
 
-
 def resnet10(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
     """ResNet-10 with optional pretrained support when `spatial_dims` is 3.
 

monai/networks/utils.py

@@ -36,6 +36,8 @@
 onnx, _ = optional_import("onnx")
 onnxreference, _ = optional_import("onnx.reference")
 onnxruntime, _ = optional_import("onnxruntime")
+hf_hub_download, _ = optional_import("huggingface_hub", name = "hf_hub_download")
+EntryNotFoundError, _ = optional_import("huggingface_hub.utils._errors", name = "EntryNotFoundError")
 
 __all__ = [
     "one_hot",
@@ -1164,3 +1166,48 @@ def freeze_layers(model: nn.Module, freeze_vars=None, exclude_vars=None):
                 warnings.warn(f"The exclude_vars includes {param}, but requires_grad is False, change it to True.")
 
     logger.info(f"{len(frozen_keys)} of {len(src_dict)} variables frozen.")
+
+def get_pretrained_resnet_medicalnet(resnet_depth: int, device: torch.device = torch.device("cpu"), datasets23: bool = True):
+    """
+    Donwlad resnet pretrained weights from https://huggingface.co/TencentMedicalNet
+
+    Args:
+        resnet_depth: depth of the pretrained model. Supported values are 10, 18, 34, 50, 101, 152 and 200
+        device: device on which the returned state dict will be loaded.
+        datasets23: if True, get the weights trained on more datasets (23). Not all depths are available. If not, standard weights are returned.
+
+    Returns:
+        Pretrained state dict
+
+    Raises:
+        huggingface_hub.utils._errors.EntryNotFoundError: if pretrained weights are not found on huggingface hub
+        NotImplementedError: if `resnet_depth` is not supported
+    """
+
+    MEDICALNET_HUGGINGFACE_REPO_BASENAME = "TencentMedicalNet/MedicalNet-Resnet"
+    MEDICALNET_HUGGINGFACE_FILES_BASENAME = "resnet_"
+    SUPPORTED_DEPTH = [10, 18, 34, 50, 101, 152, 200]
+
+    logger.info(f"Loading MedicalNet pretrained model from https://huggingface.co/{MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}")
+
+    if resnet_depth in SUPPORTED_DEPTH:
+        filename = f"{MEDICALNET_HUGGINGFACE_FILES_BASENAME}{resnet_depth}.pth" if not datasets23 else f"{MEDICALNET_HUGGINGFACE_FILES_BASENAME}{resnet_depth}_23dataset.pth"
+        try:
+            pretrained_path = hf_hub_download(
+                repo_id=f"{MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}",
+                filename=filename)
+        except Exception:
+            if datasets23:
+                logger.info(f"{filename} not available for resnet{resnet_depth}")
+                filename = f"{MEDICALNET_HUGGINGFACE_FILES_BASENAME}{resnet_depth}.pth"
+                logger.info(f"Trying with {filename}")
+                pretrained_path = hf_hub_download(
+                    repo_id=f"{MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}",
+                    filename=filename)
+            else:
+                raise EntryNotFoundError(f"{filename} not found on {MEDICALNET_HUGGINGFACE_REPO_BASENAME}{resnet_depth}")
+        checkpoint = torch.load(pretrained_path, map_location=device)
+    else:
+        raise NotImplementedError("Supported resnet_depth are: [10, 18, 34, 50, 101, 152, 200]")
+    logger.info(f"{filename} downloaded")
+    return checkpoint.get("state_dict")

tests/test_resnet.py

@@ -23,13 +23,19 @@
 from monai.utils import optional_import
 from tests.utils import test_script_save
 
+copy, _ = optional_import("copy")
+os, _ = optional_import("os")
+re, _ = optional_import("re")
+
 if TYPE_CHECKING:
     import torchvision
 
     has_torchvision = True
 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 +165,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 +189,31 @@ 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(net.state_dict(), tmp_ckpt_filename)
+
+        cp_input_param = copy.copy(input_param)
+        # Custom pretrained weights
+        cp_input_param["pretrained"] = tmp_ckpt_filename
+        pretrained_net = model(**cp_input_param)
+        assert str(net.state_dict()) == str(pretrained_net.state_dict())
+
+        # True flag
+        cp_input_param["pretrained"] = True
+        int(re.search(r"resnet(\d+)", model.__name__).group(1))
+        model(**cp_input_param)
+        if input_param.get("spatial_dims", 3) == 3:
+            model(**cp_input_param)
+        else:
+            with self.assertRaises(NotImplementedError):
+                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)