Torchvision pretrain fix (#8563)

Fixes #8552.

### Description

This updates how pretrained model weights are loaded through
Torchvision. This may not preserve historical results if the weights
being loaded are now different since the "DEFAULT" weights may not be
the weights loaded when using the `pretrained=True` argument. I tried to
preserved behaviour as indicated in the Torchvision source code where
possible.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Eric Kerfoot <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ericspod

monai/networks/blocks/fcn.py

@@ -123,7 +123,9 @@ def __init__(
         self.upsample_mode = upsample_mode
         self.conv2d_type = conv2d_type
         self.out_channels = out_channels
-        resnet = models.resnet50(pretrained=pretrained, progress=progress)
+        resnet = models.resnet50(
+            progress=progress, weights=models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None
+        )
 
         self.conv1 = resnet.conv1
         self.bn0 = resnet.bn1

monai/networks/nets/milmodel.py

@@ -16,7 +16,7 @@
 import torch
 import torch.nn as nn
 
-from monai.utils.module import optional_import
+from monai.utils import optional_import
 
 models, _ = optional_import("torchvision.models")
 
@@ -48,7 +48,6 @@ class MILModel(nn.Module):
             Defaults to ``None`` (necessary only when using a custom backbone)
         trans_blocks: number of the blocks in `TransformEncoder` layer.
         trans_dropout: dropout rate in `TransformEncoder` layer.
-
     """
 
     def __init__(
@@ -74,7 +73,7 @@ def __init__(
         self.transformer: nn.Module | None = None
 
         if backbone is None:
-            net = models.resnet50(pretrained=pretrained)
+            net = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None)
             nfc = net.fc.in_features  # save the number of final features
             net.fc = torch.nn.Identity()  # remove final linear layer
 
@@ -99,7 +98,7 @@ def hook(module, input, output):
             torch_model = getattr(models, backbone, None)
             if torch_model is None:
                 raise ValueError("Unknown torch vision model" + str(backbone))
-            net = torch_model(pretrained=pretrained)
+            net = torch_model(weights="DEFAULT" if pretrained else None)
 
             if getattr(net, "fc", None) is not None:
                 nfc = net.fc.in_features  # save the number of final features

monai/networks/nets/torchvision_fc.py

@@ -112,12 +112,11 @@ def __init__(
         weights=None,
         **kwargs,
     ):
-        if weights is not None:
-            model = getattr(models, model_name)(weights=weights, **kwargs)
-        elif pretrained:
-            model = getattr(models, model_name)(weights="DEFAULT", **kwargs)
-        else:
-            model = getattr(models, model_name)(weights=None, **kwargs)
+        # if pretrained is False, weights is a weight tensor or None for no pretrained loading
+        if pretrained and weights is None:
+            weights = "DEFAULT"
+
+        model = getattr(models, model_name)(weights=weights, **kwargs)
 
         super().__init__(
             model=model,

monai/transforms/io/array.py

@@ -43,9 +43,9 @@
 from monai.data.utils import is_no_channel
 from monai.transforms.transform import Transform
 from monai.transforms.utility.array import EnsureChannelFirst
-from monai.utils import GridSamplePadMode
-from monai.utils import ImageMetaKey as Key
 from monai.utils import (
+    GridSamplePadMode,
+    ImageMetaKey,
     MetaKeys,
     OptionalImportError,
     convert_to_dst_type,
@@ -293,7 +293,8 @@ def __call__(self, filename: Sequence[PathLike] | PathLike, reader: ImageReader
         # make sure all elements in metadata are little endian
         meta_data = switch_endianness(meta_data, "<")
 
-        meta_data[Key.FILENAME_OR_OBJ] = f"{ensure_tuple(filename)[0]}"  # Path obj should be strings for data loader
+        # Path obj should be strings for data loader
+        meta_data[ImageMetaKey.FILENAME_OR_OBJ] = f"{ensure_tuple(filename)[0]}"
         img = MetaTensor.ensure_torch_and_prune_meta(
             img_array, meta_data, self.simple_keys, pattern=self.pattern, sep=self.sep
         )
@@ -548,7 +549,7 @@ def __call__(self, img: NdarrayOrTensor):
                 "Missing 'saved_to' key in metadata. Check SaveImage argument 'savepath_in_metadict' is True."
             )
 
-        input_path = meta_data[Key.FILENAME_OR_OBJ]
+        input_path = meta_data[ImageMetaKey.FILENAME_OR_OBJ]
         output_path = meta_data[MetaKeys.SAVED_TO]
         log_data = {"input": input_path, "output": output_path}
 

requirements-dev.txt

@@ -18,7 +18,7 @@ pep8-naming
 pycodestyle
 pyflakes
 black>=25.1.0
-isort>=5.1, <6.0
+isort>=5.1, !=6.0.0
 ruff
 pytype>=2020.6.1, <=2024.4.11; platform_system != "Windows"
 types-setuptools

tests/networks/nets/test_densenet.py

@@ -96,7 +96,7 @@ def test_pretrain_consistency(self, model, input_param, input_shape):
             net = model(**input_param).to(device)
         with eval_mode(net):
             result = net.features.forward(example)
-        torchvision_net = torchvision.models.densenet121(pretrained=True).to(device)
+        torchvision_net = torchvision.models.densenet121(weights="DEFAULT").to(device)
         with eval_mode(torchvision_net):
             expected_result = torchvision_net.features.forward(example)
         self.assertTrue(torch.all(result == expected_result))

tests/networks/nets/test_milmodel.py

@@ -44,13 +44,13 @@
     TEST_CASE_MILMODEL.append(test_case)
 
 # torchvision backbone
-TEST_CASE_MILMODEL.append(
-    [{"num_classes": 5, "backbone": "resnet18", "pretrained": False}, (2, 2, 3, 512, 512), (2, 5)]
-)
-TEST_CASE_MILMODEL.append([{"num_classes": 5, "backbone": "resnet18", "pretrained": True}, (2, 2, 3, 512, 512), (2, 5)])
+for pretrained in [True, False]:
+    TEST_CASE_MILMODEL.append(
+        [{"num_classes": 5, "backbone": "resnet18", "pretrained": pretrained}, (2, 2, 3, 512, 512), (2, 5)]
+    )
 
 # custom backbone
-backbone = models.densenet121(pretrained=False)
+backbone = models.densenet121()
 backbone_nfeatures = backbone.classifier.in_features
 backbone.classifier = torch.nn.Identity()
 TEST_CASE_MILMODEL.append(