Merge branch 'main' into bugfix/autocontrast_equal_minmax

Author: datumbox

.circleci/config.yml

@@ -278,8 +278,9 @@ jobs:
           background: true
           command: |
             sudo apt update -qy && sudo apt install -qy parallel wget
+            mkdir -p ~/.cache/torch/hub/checkpoints
             python scripts/collect_model_urls.py torchvision/prototype/models \
-                | parallel -j0 wget --no-verbose -P ~/.cache/torch/hub/checkpoints {}
+                | parallel -j0 'wget --no-verbose -O ~/.cache/torch/hub/checkpoints/`basename {}` {}\?source=ci'
       - run:
           name: Install torchvision
           command: pip install --user --progress-bar off --no-build-isolation .

.circleci/config.yml.in

@@ -4,7 +4,7 @@
 import time
 
 import torch
-import torch.quantization
+import torch.ao.quantization
 import torch.utils.data
 import torchvision
 import utils
@@ -62,8 +62,8 @@ def main(args):
 
     if not (args.test_only or args.post_training_quantize):
         model.fuse_model()
-        model.qconfig = torch.quantization.get_default_qat_qconfig(args.backend)
-        torch.quantization.prepare_qat(model, inplace=True)
+        model.qconfig = torch.ao.quantization.get_default_qat_qconfig(args.backend)
+        torch.ao.quantization.prepare_qat(model, inplace=True)
 
         if args.distributed and args.sync_bn:
             model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
@@ -96,12 +96,12 @@ def main(args):
         )
         model.eval()
         model.fuse_model()
-        model.qconfig = torch.quantization.get_default_qconfig(args.backend)
-        torch.quantization.prepare(model, inplace=True)
+        model.qconfig = torch.ao.quantization.get_default_qconfig(args.backend)
+        torch.ao.quantization.prepare(model, inplace=True)
         # Calibrate first
         print("Calibrating")
         evaluate(model, criterion, data_loader_calibration, device=device, print_freq=1)
-        torch.quantization.convert(model, inplace=True)
+        torch.ao.quantization.convert(model, inplace=True)
         if args.output_dir:
             print("Saving quantized model")
             if utils.is_main_process():
@@ -114,8 +114,8 @@ def main(args):
         evaluate(model, criterion, data_loader_test, device=device)
         return
 
-    model.apply(torch.quantization.enable_observer)
-    model.apply(torch.quantization.enable_fake_quant)
+    model.apply(torch.ao.quantization.enable_observer)
+    model.apply(torch.ao.quantization.enable_fake_quant)
     start_time = time.time()
     for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
@@ -126,7 +126,7 @@ def main(args):
         with torch.inference_mode():
             if epoch >= args.num_observer_update_epochs:
                 print("Disabling observer for subseq epochs, epoch = ", epoch)
-                model.apply(torch.quantization.disable_observer)
+                model.apply(torch.ao.quantization.disable_observer)
             if epoch >= args.num_batch_norm_update_epochs:
                 print("Freezing BN for subseq epochs, epoch = ", epoch)
                 model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
@@ -136,7 +136,7 @@ def main(args):
             quantized_eval_model = copy.deepcopy(model_without_ddp)
             quantized_eval_model.eval()
             quantized_eval_model.to(torch.device("cpu"))
-            torch.quantization.convert(quantized_eval_model, inplace=True)
+            torch.ao.quantization.convert(quantized_eval_model, inplace=True)
 
             print("Evaluate Quantized model")
             evaluate(quantized_eval_model, criterion, data_loader_test, device=torch.device("cpu"))

references/classification/train_quantization.py

@@ -345,8 +345,8 @@ def store_model_weights(model, checkpoint_path, checkpoint_key="model", strict=T
         # Quantized Classification
         model = M.quantization.mobilenet_v3_large(pretrained=False, quantize=False)
         model.fuse_model()
-        model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
-        _ = torch.quantization.prepare_qat(model, inplace=True)
+        model.qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack')
+        _ = torch.ao.quantization.prepare_qat(model, inplace=True)
         print(store_model_weights(model, './qat.pth'))
 
         # Object Detection

references/classification/utils.py

@@ -795,6 +795,36 @@ def test_solarize2(device, dtype, config, channels):
     )
 
 
+@pytest.mark.parametrize("device", cpu_and_gpu())
+@pytest.mark.parametrize("threshold", [0.0, 0.25, 0.5, 0.75, 1.0])
+def test_solarize_threshold1_bound(threshold, device):
+    img = torch.rand((3, 12, 23)).to(device)
+    F_t.solarize(img, threshold)
+
+
+@pytest.mark.parametrize("device", cpu_and_gpu())
+@pytest.mark.parametrize("threshold", [1.5])
+def test_solarize_threshold1_upper_bound(threshold, device):
+    img = torch.rand((3, 12, 23)).to(device)
+    with pytest.raises(TypeError, match="Threshold should be less than bound of img."):
+        F_t.solarize(img, threshold)
+
+
+@pytest.mark.parametrize("device", cpu_and_gpu())
+@pytest.mark.parametrize("threshold", [0, 64, 128, 192, 255])
+def test_solarize_threshold2_bound(threshold, device):
+    img = torch.randint(0, 256, (3, 12, 23)).to(device)
+    F_t.solarize(img, threshold)
+
+
+@pytest.mark.parametrize("device", cpu_and_gpu())
+@pytest.mark.parametrize("threshold", [260])
+def test_solarize_threshold2_upper_bound(threshold, device):
+    img = torch.randint(0, 256, (3, 12, 23)).to(device)
+    with pytest.raises(TypeError, match="Threshold should be less than bound of img."):
+        F_t.solarize(img, threshold)
+
+
 @pytest.mark.parametrize("device", cpu_and_gpu())
 @pytest.mark.parametrize("dtype", (None, torch.float32, torch.float64))
 @pytest.mark.parametrize("config", [{"sharpness_factor": f} for f in [0.2, 0.5, 1.0, 1.5, 2.0]])

test/test_functional_tensor.py

@@ -781,19 +781,19 @@ def test_quantized_classification_model(model_fn):
         model = model_fn(**kwargs)
         if eval_mode:
             model.eval()
-            model.qconfig = torch.quantization.default_qconfig
+            model.qconfig = torch.ao.quantization.default_qconfig
         else:
             model.train()
-            model.qconfig = torch.quantization.default_qat_qconfig
+            model.qconfig = torch.ao.quantization.default_qat_qconfig
 
         model.fuse_model()
         if eval_mode:
-            torch.quantization.prepare(model, inplace=True)
+            torch.ao.quantization.prepare(model, inplace=True)
         else:
-            torch.quantization.prepare_qat(model, inplace=True)
+            torch.ao.quantization.prepare_qat(model, inplace=True)
             model.eval()
 
-        torch.quantization.convert(model, inplace=True)
+        torch.ao.quantization.convert(model, inplace=True)
 
     try:
         torch.jit.script(model)

test/test_models.py

@@ -18,6 +18,12 @@ def _get_original_model(model_fn):
     return module.__dict__[model_fn.__name__]
 
 
+def _get_parent_module(model_fn):
+    parent_module_name = ".".join(model_fn.__module__.split(".")[:-1])
+    module = importlib.import_module(parent_module_name)
+    return module
+
+
 def _build_model(fn, **kwargs):
     try:
         model = fn(**kwargs)
@@ -29,27 +35,42 @@ def _build_model(fn, **kwargs):
     return model.eval()
 
 
-def get_models_with_module_names(module):
-    module_name = module.__name__.split(".")[-1]
-    return [(fn, module_name) for fn in TM.get_models_from_module(module)]
-
-
 @pytest.mark.parametrize(
     "model_fn, name, weight",
     [
-        (models.resnet50, "ImageNet1K_RefV1", models.ResNet50Weights.ImageNet1K_RefV1),
-        (models.resnet50, "default", models.ResNet50Weights.ImageNet1K_RefV2),
+        (models.resnet50, "ImageNet1K_V1", models.ResNet50_Weights.ImageNet1K_V1),
+        (models.resnet50, "default", models.ResNet50_Weights.ImageNet1K_V2),
+        (
+            models.quantization.resnet50,
+            "default",
+            models.quantization.ResNet50_QuantizedWeights.ImageNet1K_FBGEMM_V2,
+        ),
         (
             models.quantization.resnet50,
-            "ImageNet1K_FBGEMM_RefV1",
-            models.quantization.QuantizedResNet50Weights.ImageNet1K_FBGEMM_RefV1,
+            "ImageNet1K_FBGEMM_V1",
+            models.quantization.ResNet50_QuantizedWeights.ImageNet1K_FBGEMM_V1,
         ),
     ],
 )
 def test_get_weight(model_fn, name, weight):
     assert models._api.get_weight(model_fn, name) == weight
 
 
+@pytest.mark.parametrize(
+    "model_fn",
+    TM.get_models_from_module(models)
+    + TM.get_models_from_module(models.detection)
+    + TM.get_models_from_module(models.quantization)
+    + TM.get_models_from_module(models.segmentation)
+    + TM.get_models_from_module(models.video),
+)
+def test_naming_conventions(model_fn):
+    model_name = model_fn.__name__
+    module = _get_parent_module(model_fn)
+    weights_name = "_QuantizedWeights" if module.__name__.split(".")[-1] == "quantization" else "_Weights"
+    assert model_name in set(x.replace(weights_name, "").lower() for x in module.__dict__ if x.endswith(weights_name))
+
+
 @pytest.mark.parametrize("model_fn", TM.get_models_from_module(models))
 @pytest.mark.parametrize("dev", cpu_and_gpu())
 @run_if_test_with_prototype
@@ -85,16 +106,16 @@ def test_video_model(model_fn, dev):
 
 
 @pytest.mark.parametrize(
-    "model_fn, module_name",
-    get_models_with_module_names(models)
-    + get_models_with_module_names(models.detection)
-    + get_models_with_module_names(models.quantization)
-    + get_models_with_module_names(models.segmentation)
-    + get_models_with_module_names(models.video),
+    "model_fn",
+    TM.get_models_from_module(models)
+    + TM.get_models_from_module(models.detection)
+    + TM.get_models_from_module(models.quantization)
+    + TM.get_models_from_module(models.segmentation)
+    + TM.get_models_from_module(models.video),
 )
 @pytest.mark.parametrize("dev", cpu_and_gpu())
 @run_if_test_with_prototype
-def test_old_vs_new_factory(model_fn, module_name, dev):
+def test_old_vs_new_factory(model_fn, dev):
     defaults = {
         "models": {
             "input_shape": (1, 3, 224, 224),
@@ -114,6 +135,7 @@ def test_old_vs_new_factory(model_fn, module_name, dev):
         },
     }
     model_name = model_fn.__name__
+    module_name = model_fn.__module__.split(".")[-2]
     kwargs = {"pretrained": True, **defaults[module_name], **TM._model_params.get(model_name, {})}
     input_shape = kwargs.pop("input_shape")
     kwargs.pop("num_classes", None)  # ignore this as it's an incompatible speed optimization for pre-trained models

test/test_prototype_models.py

@@ -31,7 +31,7 @@ def forward(self, x: Tensor) -> Tensor:
         return x
 
     def fuse_model(self) -> None:
-        torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
+        torch.ao.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
 
 class QuantizableInception(Inception):
@@ -74,8 +74,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(  # type: ignore[misc]
             blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux], *args, **kwargs
         )
-        self.quant = torch.quantization.QuantStub()
-        self.dequant = torch.quantization.DeQuantStub()
+        self.quant = torch.ao.quantization.QuantStub()
+        self.dequant = torch.ao.quantization.DeQuantStub()
 
     def forward(self, x: Tensor) -> GoogLeNetOutputs:
         x = self._transform_input(x)

torchvision/models/quantization/googlenet.py

@@ -36,7 +36,7 @@ def forward(self, x: Tensor) -> Tensor:
         return x
 
     def fuse_model(self) -> None:
-        torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
+        torch.ao.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
 
 class QuantizableInceptionA(inception_module.InceptionA):
@@ -144,8 +144,8 @@ def __init__(
                 QuantizableInceptionAux,
             ],
         )
-        self.quant = torch.quantization.QuantStub()
-        self.dequant = torch.quantization.DeQuantStub()
+        self.quant = torch.ao.quantization.QuantStub()
+        self.dequant = torch.ao.quantization.DeQuantStub()
 
     def forward(self, x: Tensor) -> InceptionOutputs:
         x = self._transform_input(x)

torchvision/models/quantization/inception.py

@@ -2,7 +2,7 @@
 
 from torch import Tensor
 from torch import nn
-from torch.quantization import QuantStub, DeQuantStub, fuse_modules
+from torch.ao.quantization import QuantStub, DeQuantStub, fuse_modules
 from torchvision.models.mobilenetv2 import InvertedResidual, MobileNetV2, model_urls
 
 from ..._internally_replaced_utils import load_state_dict_from_url

torchvision/models/quantization/mobilenetv2.py

@@ -2,7 +2,7 @@
 
 import torch
 from torch import nn, Tensor
-from torch.quantization import QuantStub, DeQuantStub, fuse_modules
+from torch.ao.quantization import QuantStub, DeQuantStub, fuse_modules
 
 from ..._internally_replaced_utils import load_state_dict_from_url
 from ...ops.misc import ConvNormActivation, SqueezeExcitation
@@ -136,13 +136,13 @@ def _mobilenet_v3_model(
         backend = "qnnpack"
 
         model.fuse_model()
-        model.qconfig = torch.quantization.get_default_qat_qconfig(backend)
-        torch.quantization.prepare_qat(model, inplace=True)
+        model.qconfig = torch.ao.quantization.get_default_qat_qconfig(backend)
+        torch.ao.quantization.prepare_qat(model, inplace=True)
 
         if pretrained:
             _load_weights(arch, model, quant_model_urls.get(arch + "_" + backend, None), progress)
 
-        torch.quantization.convert(model, inplace=True)
+        torch.ao.quantization.convert(model, inplace=True)
         model.eval()
     else:
         if pretrained:

torchvision/models/quantization/mobilenetv3.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
-from torch.quantization import fuse_modules
+from torch.ao.quantization import fuse_modules
 from torchvision.models.resnet import Bottleneck, BasicBlock, ResNet, model_urls
 
 from ..._internally_replaced_utils import load_state_dict_from_url
@@ -42,9 +42,9 @@ def forward(self, x: Tensor) -> Tensor:
         return out
 
     def fuse_model(self) -> None:
-        torch.quantization.fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True)
+        torch.ao.quantization.fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True)
         if self.downsample:
-            torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
+            torch.ao.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
 
 
 class QuantizableBottleneck(Bottleneck):
@@ -75,15 +75,15 @@ def forward(self, x: Tensor) -> Tensor:
     def fuse_model(self) -> None:
         fuse_modules(self, [["conv1", "bn1", "relu1"], ["conv2", "bn2", "relu2"], ["conv3", "bn3"]], inplace=True)
         if self.downsample:
-            torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
+            torch.ao.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
 
 
 class QuantizableResNet(ResNet):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
-        self.quant = torch.quantization.QuantStub()
-        self.dequant = torch.quantization.DeQuantStub()
+        self.quant = torch.ao.quantization.QuantStub()
+        self.dequant = torch.ao.quantization.DeQuantStub()
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)

torchvision/models/quantization/resnet.py

@@ -41,8 +41,8 @@ class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
     # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, inverted_residual=QuantizableInvertedResidual, **kwargs)  # type: ignore[misc]
-        self.quant = torch.quantization.QuantStub()
-        self.dequant = torch.quantization.DeQuantStub()
+        self.quant = torch.ao.quantization.QuantStub()
+        self.dequant = torch.ao.quantization.DeQuantStub()
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)
@@ -60,12 +60,12 @@ def fuse_model(self) -> None:
 
         for name, m in self._modules.items():
             if name in ["conv1", "conv5"]:
-                torch.quantization.fuse_modules(m, [["0", "1", "2"]], inplace=True)
+                torch.ao.quantization.fuse_modules(m, [["0", "1", "2"]], inplace=True)
         for m in self.modules():
             if type(m) is QuantizableInvertedResidual:
                 if len(m.branch1._modules.items()) > 0:
-                    torch.quantization.fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True)
-                torch.quantization.fuse_modules(
+                    torch.ao.quantization.fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True)
+                torch.ao.quantization.fuse_modules(
                     m.branch2,
                     [["0", "1", "2"], ["3", "4"], ["5", "6", "7"]],
                     inplace=True,