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 <17726042+ericspod@users.noreply.github.com> 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(