[TorchFix] Make pretrained rule work when a model imported from a submodule (#4613)

Update the rule for `pretrained` torchvision models parameter to work
when a model imported from a submodule, like `models.resnet.resnet101`
instead of `models.resnet101` (which is the same model).

Author: kit1980

tests/fixtures/vision/checker/pretrained.py

@@ -17,3 +17,7 @@
 from torchvision.models import ResNet50_Weights
 torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
 torchvision.models.resnet50(weights=None)
+
+# Make sure no false positives on non-model functions
+from torchvision.models.vgg import make_layers, cfgs
+make_layers(cfgs['D'], False)

tests/fixtures/vision/codemod/pretrained_submodule_import.py

@@ -0,0 +1,14 @@
+# Check that the codemod works when a model
+# imported from a submodule
+# i.e. from torchvision.models.resnet instead of
+# directly from torchvision.models.
+
+from torchvision import models
+backbone = models.resnet.resnet101(pretrained=True, replace_stride_with_dilation=[False, True, True])
+
+from torchvision.models import resnet
+backbone = resnet.resnet101(pretrained=True, replace_stride_with_dilation=[False, True, True])
+
+from torchvision.models.resnet import resnet152
+resnet152(pretrained=True)
+resnet152(pretrained=False)

tests/fixtures/vision/codemod/pretrained_submodule_import.py.out

@@ -0,0 +1,14 @@
+# Check that the codemod works when a model
+# imported from a submodule
+# i.e. from torchvision.models.resnet instead of
+# directly from torchvision.models.
+
+from torchvision import models
+backbone = models.resnet.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V1, replace_stride_with_dilation=[False, True, True])
+
+from torchvision.models import resnet
+backbone = resnet.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V1, replace_stride_with_dilation=[False, True, True])
+
+from torchvision.models.resnet import resnet152
+resnet152(weights=models.ResNet152_Weights.IMAGENET1K_V1)
+resnet152(weights=None)

torchfix/visitors/vision/__init__.py

@@ -150,6 +150,27 @@ class TorchVisionDeprecatedPretrainedVisitor(TorchVisitor):
     }
     # fmt: on
 
+    # The same model can be imported from torchvision.models directly,
+    # or from a submodule like torchvision.models.resnet.
+    MODEL_SUBMODULES = (
+        "alexnet",
+        "convnext",
+        "densenet",
+        "efficientnet",
+        "googlenet",
+        "inception",
+        "mnasnet",
+        "mobilenet",
+        "regnet",
+        "resnet",
+        "shufflenetv2",
+        "squeezenet",
+        "vgg",
+        "vision_transformer",
+        "swin_transformer",
+        "maxvit",
+    )
+
     def visit_Call(self, node):
         def _new_arg_and_import(
             old_arg: cst.Arg, is_backbone: bool
@@ -185,8 +206,14 @@ def _new_arg_and_import(
             return
         if qualified_name.startswith("torchvision.models"):
             model_name = qualified_name[len("torchvision.models") + 1 :]
-            message = None
+            for submodule in self.MODEL_SUBMODULES:
+                if model_name.startswith(submodule + "."):
+                    model_name = model_name[len(submodule) + 1 :]
 
+            if (model_name, "pretrained") not in self.MODEL_WEIGHTS:
+                return
+
+            message = None
             pretrained_arg = self.get_specific_arg(node, "pretrained", 0)
             if pretrained_arg is not None:
                 message = "Parameter `pretrained` is deprecated, please use `weights` instead."