Merge branch 'main' into add_flow_datasets

Author: NicolasHug

docs/source/models.rst

@@ -479,7 +479,8 @@ Model                             Acc@1          Acc@5
 ================================  =============  =============
 MobileNet V2                      71.658         90.150
 MobileNet V3 Large                73.004         90.858
-ShuffleNet V2                     68.360         87.582
+ShuffleNet V2 x1.0                68.360         87.582
+ShuffleNet V2 x0.5                57.972         79.780
 ResNet 18                         69.494         88.882
 ResNet 50                         75.920         92.814
 ResNext 101 32x8d                 78.986         94.480

references/classification/README.md

@@ -168,7 +168,7 @@ For all post training quantized models, the settings are:
 ```
 python train_quantization.py --device='cpu' --post-training-quantize --backend='fbgemm' --model='$MODEL'
 ```
-Here `$MODEL` is one of `googlenet`, `inception_v3`, `resnet18`, `resnet50`, `resnext101_32x8d` and `shufflenet_v2_x1_0`.
+Here `$MODEL` is one of `googlenet`, `inception_v3`, `resnet18`, `resnet50`, `resnext101_32x8d`, `shufflenet_v2_x0_5` and `shufflenet_v2_x1_0`.
 
 ### QAT MobileNetV2
 

test/test_transforms.py

@@ -1364,7 +1364,7 @@ def test_to_grayscale():
 @pytest.mark.parametrize("p", (0, 1))
 def test_random_apply(p, seed):
     torch.manual_seed(seed)
-    random_apply_transform = transforms.RandomApply([transforms.RandomRotation((1, 45))], p=p)
+    random_apply_transform = transforms.RandomApply([transforms.RandomRotation((45, 50))], p=p)
     img = transforms.ToPILImage()(torch.rand(3, 30, 40))
     out = random_apply_transform(img)
     if p == 0:
@@ -1384,7 +1384,7 @@ def test_random_choice(proba_passthrough, seed):
     random_choice_transform = transforms.RandomChoice(
         [
             lambda x: x,  # passthrough
-            transforms.RandomRotation((1, 45)),
+            transforms.RandomRotation((45, 50)),
         ],
         p=[proba_passthrough, 1 - proba_passthrough],
     )

torchvision/models/quantization/shufflenetv2.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Optional
 
 import torch
 import torch.nn as nn
@@ -12,15 +12,11 @@
     "QuantizableShuffleNetV2",
     "shufflenet_v2_x0_5",
     "shufflenet_v2_x1_0",
-    "shufflenet_v2_x1_5",
-    "shufflenet_v2_x2_0",
 ]
 
 quant_model_urls = {
-    "shufflenetv2_x0.5_fbgemm": None,
+    "shufflenetv2_x0.5_fbgemm": "https://download.pytorch.org/models/quantized/shufflenetv2_x0.5_fbgemm-00845098.pth",
     "shufflenetv2_x1.0_fbgemm": "https://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-db332c57.pth",
-    "shufflenetv2_x1.5_fbgemm": None,
-    "shufflenetv2_x2.0_fbgemm": None,
 }
 
 
@@ -96,6 +92,7 @@ def _shufflenetv2(
         assert pretrained in [True, False]
 
     if pretrained:
+        model_url: Optional[str] = None
         if quantize:
             model_url = quant_model_urls[arch + "_" + backend]
         else:
@@ -147,45 +144,3 @@ def shufflenet_v2_x1_0(
     return _shufflenetv2(
         "shufflenetv2_x1.0", pretrained, progress, quantize, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs
     )
-
-
-def shufflenet_v2_x1_5(
-    pretrained: bool = False,
-    progress: bool = True,
-    quantize: bool = False,
-    **kwargs: Any,
-) -> QuantizableShuffleNetV2:
-    """
-    Constructs a ShuffleNetV2 with 1.5x output channels, as described in
-    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
-    <https://arxiv.org/abs/1807.11164>`_.
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-        progress (bool): If True, displays a progress bar of the download to stderr
-        quantize (bool): If True, return a quantized version of the model
-    """
-    return _shufflenetv2(
-        "shufflenetv2_x1.5", pretrained, progress, quantize, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs
-    )
-
-
-def shufflenet_v2_x2_0(
-    pretrained: bool = False,
-    progress: bool = True,
-    quantize: bool = False,
-    **kwargs: Any,
-) -> QuantizableShuffleNetV2:
-    """
-    Constructs a ShuffleNetV2 with 2.0x output channels, as described in
-    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
-    <https://arxiv.org/abs/1807.11164>`_.
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-        progress (bool): If True, displays a progress bar of the download to stderr
-        quantize (bool): If True, return a quantized version of the model
-    """
-    return _shufflenetv2(
-        "shufflenetv2_x2.0", pretrained, progress, quantize, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs
-    )

torchvision/prototype/models/quantization/__init__.py

@@ -1,3 +1,4 @@
 from .googlenet import *
 from .inception import *
 from .resnet import *
+from .shufflenetv2 import *

torchvision/prototype/models/quantization/shufflenetv2.py

@@ -0,0 +1,135 @@
+import warnings
+from functools import partial
+from typing import Any, List, Optional, Union
+
+from torchvision.transforms.functional import InterpolationMode
+
+from ....models.quantization.shufflenetv2 import (
+    QuantizableShuffleNetV2,
+    _replace_relu,
+    quantize_model,
+)
+from ...transforms.presets import ImageNetEval
+from .._api import Weights, WeightEntry
+from .._meta import _IMAGENET_CATEGORIES
+from ..shufflenetv2 import ShuffleNetV2_x0_5Weights, ShuffleNetV2_x1_0Weights
+
+
+__all__ = [
+    "QuantizableShuffleNetV2",
+    "QuantizedShuffleNetV2_x0_5Weights",
+    "QuantizedShuffleNetV2_x1_0Weights",
+    "shufflenet_v2_x0_5",
+    "shufflenet_v2_x1_0",
+]
+
+
+def _shufflenetv2(
+    stages_repeats: List[int],
+    stages_out_channels: List[int],
+    weights: Optional[Weights],
+    progress: bool,
+    quantize: bool,
+    **kwargs: Any,
+) -> QuantizableShuffleNetV2:
+    if weights is not None:
+        kwargs["num_classes"] = len(weights.meta["categories"])
+        if "backend" in weights.meta:
+            kwargs["backend"] = weights.meta["backend"]
+    backend = kwargs.pop("backend", "fbgemm")
+
+    model = QuantizableShuffleNetV2(stages_repeats, stages_out_channels, **kwargs)
+    _replace_relu(model)
+    if quantize:
+        quantize_model(model, backend)
+
+    if weights is not None:
+        model.load_state_dict(weights.state_dict(progress=progress))
+
+    return model
+
+
+_common_meta = {
+    "size": (224, 224),
+    "categories": _IMAGENET_CATEGORIES,
+    "interpolation": InterpolationMode.BILINEAR,
+    "backend": "fbgemm",
+    "quantization": "ptq",
+    "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models",
+}
+
+
+class QuantizedShuffleNetV2_x0_5Weights(Weights):
+    ImageNet1K_FBGEMM_Community = WeightEntry(
+        url="https://download.pytorch.org/models/quantized/shufflenetv2_x0.5_fbgemm-00845098.pth",
+        transforms=partial(ImageNetEval, crop_size=224),
+        meta={
+            **_common_meta,
+            "unquantized": ShuffleNetV2_x0_5Weights.ImageNet1K_Community,
+            "acc@1": 57.972,
+            "acc@5": 79.780,
+        },
+    )
+
+
+class QuantizedShuffleNetV2_x1_0Weights(Weights):
+    ImageNet1K_FBGEMM_Community = WeightEntry(
+        url="https://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-db332c57.pth",
+        transforms=partial(ImageNetEval, crop_size=224),
+        meta={
+            **_common_meta,
+            "unquantized": ShuffleNetV2_x1_0Weights.ImageNet1K_Community,
+            "acc@1": 68.360,
+            "acc@5": 87.582,
+        },
+    )
+
+
+def shufflenet_v2_x0_5(
+    weights: Optional[Union[QuantizedShuffleNetV2_x0_5Weights, ShuffleNetV2_x0_5Weights]] = None,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any,
+) -> QuantizableShuffleNetV2:
+    if "pretrained" in kwargs:
+        warnings.warn("The argument pretrained is deprecated, please use weights instead.")
+        if kwargs.pop("pretrained"):
+            weights = (
+                QuantizedShuffleNetV2_x0_5Weights.ImageNet1K_FBGEMM_Community
+                if quantize
+                else ShuffleNetV2_x0_5Weights.ImageNet1K_Community
+            )
+        else:
+            weights = None
+
+    if quantize:
+        weights = QuantizedShuffleNetV2_x0_5Weights.verify(weights)
+    else:
+        weights = ShuffleNetV2_x0_5Weights.verify(weights)
+
+    return _shufflenetv2([4, 8, 4], [24, 48, 96, 192, 1024], weights, progress, quantize, **kwargs)
+
+
+def shufflenet_v2_x1_0(
+    weights: Optional[Union[QuantizedShuffleNetV2_x1_0Weights, ShuffleNetV2_x1_0Weights]] = None,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any,
+) -> QuantizableShuffleNetV2:
+    if "pretrained" in kwargs:
+        warnings.warn("The argument pretrained is deprecated, please use weights instead.")
+        if kwargs.pop("pretrained"):
+            weights = (
+                QuantizedShuffleNetV2_x1_0Weights.ImageNet1K_FBGEMM_Community
+                if quantize
+                else ShuffleNetV2_x1_0Weights.ImageNet1K_Community
+            )
+        else:
+            weights = None
+
+    if quantize:
+        weights = QuantizedShuffleNetV2_x1_0Weights.verify(weights)
+    else:
+        weights = ShuffleNetV2_x1_0Weights.verify(weights)
+
+    return _shufflenetv2([4, 8, 4], [24, 116, 232, 464, 1024], weights, progress, quantize, **kwargs)