Merge branch 'main' into disable-weight-loading

pmeier · web-flow · commit b6e65698822a · 2021-11-06T14:26:53.000+01:00
diff --git a/torchvision/csrc/io/decoder/stream.cpp b/torchvision/csrc/io/decoder/stream.cpp
@@ -74,15 +74,6 @@ int Stream::openCodec(std::vector<DecoderMetadata>* metadata, int num_threads) {
     }
   }
 
-  // print codec type and number of threads
-  LOG(INFO) << "Codec " << codecCtx_->codec->long_name
-            << " Codec id: " << codecCtx_->codec_id
-            << " Codec tag: " << codecCtx_->codec_tag
-            << " Codec type: " << codecCtx_->codec_type
-            << " Codec extradata: " << codecCtx_->extradata
-            << " Number of threads: " << codecCtx_->thread_count
-            << " Thread type: " << codecCtx_->thread_type;
-
   int ret;
   // Copy codec parameters from input stream to output codec context
   if ((ret = avcodec_parameters_to_context(codecCtx_, steam->codecpar)) < 0) {
diff --git a/torchvision/datasets/cifar.py b/torchvision/datasets/cifar.py
@@ -1,4 +1,3 @@
-import os
 import os.path
 import pickle
 from typing import Any, Callable, Optional, Tuple
diff --git a/torchvision/datasets/coco.py b/torchvision/datasets/coco.py
@@ -1,4 +1,3 @@
-import os
 import os.path
 from typing import Any, Callable, Optional, Tuple, List
 
diff --git a/torchvision/datasets/lsun.py b/torchvision/datasets/lsun.py
@@ -1,5 +1,4 @@
 import io
-import os
 import os.path
 import pickle
 import string
diff --git a/torchvision/datasets/semeion.py b/torchvision/datasets/semeion.py
@@ -1,4 +1,3 @@
-import os
 import os.path
 from typing import Any, Callable, Optional, Tuple
 
diff --git a/torchvision/datasets/stl10.py b/torchvision/datasets/stl10.py
@@ -1,4 +1,3 @@
-import os
 import os.path
 from typing import Any, Callable, Optional, Tuple
 
diff --git a/torchvision/datasets/svhn.py b/torchvision/datasets/svhn.py
@@ -1,4 +1,3 @@
-import os
 import os.path
 from typing import Any, Callable, Optional, Tuple
 
diff --git a/torchvision/prototype/models/_meta.py b/torchvision/prototype/models/_meta.py
@@ -1102,6 +1102,28 @@
     "toothbrush",
 ]
 
+# To be replaced with torchvision.datasets.info("coco_kp")
+_COCO_PERSON_CATEGORIES = ["no person", "person"]
+_COCO_PERSON_KEYPOINT_NAMES = [
+    "nose",
+    "left_eye",
+    "right_eye",
+    "left_ear",
+    "right_ear",
+    "left_shoulder",
+    "right_shoulder",
+    "left_elbow",
+    "right_elbow",
+    "left_wrist",
+    "right_wrist",
+    "left_hip",
+    "right_hip",
+    "left_knee",
+    "right_knee",
+    "left_ankle",
+    "right_ankle",
+]
+
 # To be replaced with torchvision.datasets.info("voc").categories
 _VOC_CATEGORIES = [
     "__background__",
diff --git a/torchvision/prototype/models/detection/__init__.py b/torchvision/prototype/models/detection/__init__.py
@@ -1,2 +1,4 @@
 from .faster_rcnn import *
+from .keypoint_rcnn import *
 from .mask_rcnn import *
+from .retinanet import *
diff --git a/torchvision/prototype/models/detection/faster_rcnn.py b/torchvision/prototype/models/detection/faster_rcnn.py
@@ -1,6 +1,8 @@
 import warnings
 from typing import Any, Optional, Union
 
+from torchvision.transforms.functional import InterpolationMode
+
 from ....models.detection.faster_rcnn import (
     _mobilenet_extractor,
     _resnet_fpn_extractor,
@@ -28,7 +30,10 @@
 ]
 
 
-_common_meta = {"categories": _COCO_CATEGORIES}
+_common_meta = {
+    "categories": _COCO_CATEGORIES,
+    "interpolation": InterpolationMode.BILINEAR,
+}
 
 
 class FasterRCNNResNet50FPNWeights(Weights):
diff --git a/torchvision/prototype/models/detection/keypoint_rcnn.py b/torchvision/prototype/models/detection/keypoint_rcnn.py
@@ -0,0 +1,92 @@
+import warnings
+from typing import Any, Optional
+
+from ....models.detection.keypoint_rcnn import (
+    _resnet_fpn_extractor,
+    _validate_trainable_layers,
+    KeypointRCNN,
+    misc_nn_ops,
+    overwrite_eps,
+)
+from ...transforms.presets import CocoEval
+from .._api import Weights, WeightEntry
+from .._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES
+from ..resnet import ResNet50Weights, resnet50
+
+
+__all__ = [
+    "KeypointRCNN",
+    "KeypointRCNNResNet50FPNWeights",
+    "keypointrcnn_resnet50_fpn",
+]
+
+
+_common_meta = {"categories": _COCO_PERSON_CATEGORIES, "keypoint_names": _COCO_PERSON_KEYPOINT_NAMES}
+
+
+class KeypointRCNNResNet50FPNWeights(Weights):
+    Coco_RefV1_Legacy = WeightEntry(
+        url="https://download.pytorch.org/models/keypointrcnn_resnet50_fpn_coco-9f466800.pth",
+        transforms=CocoEval,
+        meta={
+            **_common_meta,
+            "recipe": "https://github.com/pytorch/vision/issues/1606",
+            "box_map": 50.6,
+            "kp_map": 61.1,
+        },
+    )
+    Coco_RefV1 = WeightEntry(
+        url="https://download.pytorch.org/models/keypointrcnn_resnet50_fpn_coco-fc266e95.pth",
+        transforms=CocoEval,
+        meta={
+            **_common_meta,
+            "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#keypoint-r-cnn",
+            "box_map": 54.6,
+            "kp_map": 65.0,
+        },
+    )
+
+
+def keypointrcnn_resnet50_fpn(
+    weights: Optional[KeypointRCNNResNet50FPNWeights] = None,
+    weights_backbone: Optional[ResNet50Weights] = None,
+    progress: bool = True,
+    num_classes: int = 2,
+    num_keypoints: int = 17,
+    trainable_backbone_layers: Optional[int] = None,
+    **kwargs: Any,
+) -> KeypointRCNN:
+    if "pretrained" in kwargs:
+        warnings.warn("The argument pretrained is deprecated, please use weights instead.")
+        pretrained = kwargs.pop("pretrained")
+        if type(pretrained) == str and pretrained == "legacy":
+            weights = KeypointRCNNResNet50FPNWeights.Coco_RefV1_Legacy
+        elif type(pretrained) == bool and pretrained:
+            weights = KeypointRCNNResNet50FPNWeights.Coco_RefV1
+        else:
+            weights = None
+    weights = KeypointRCNNResNet50FPNWeights.verify(weights)
+    if "pretrained_backbone" in kwargs:
+        warnings.warn("The argument pretrained_backbone is deprecated, please use weights_backbone instead.")
+        weights_backbone = ResNet50Weights.ImageNet1K_RefV1 if kwargs.pop("pretrained_backbone") else None
+    weights_backbone = ResNet50Weights.verify(weights_backbone)
+
+    if weights is not None:
+        weights_backbone = None
+        num_classes = len(weights.meta["categories"])
+        num_keypoints = len(weights.meta["keypoint_names"])
+
+    trainable_backbone_layers = _validate_trainable_layers(
+        weights is not None or weights_backbone is not None, trainable_backbone_layers, 5, 3
+    )
+
+    backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=misc_nn_ops.FrozenBatchNorm2d)
+    backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
+    model = KeypointRCNN(backbone, num_classes, num_keypoints=num_keypoints, **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.state_dict(progress=progress))
+        if weights == KeypointRCNNResNet50FPNWeights.Coco_RefV1:
+            overwrite_eps(model, 0.0)
+
+    return model
diff --git a/torchvision/prototype/models/detection/mask_rcnn.py b/torchvision/prototype/models/detection/mask_rcnn.py
@@ -1,6 +1,8 @@
 import warnings
 from typing import Any, Optional
 
+from torchvision.transforms.functional import InterpolationMode
+
 from ....models.detection.mask_rcnn import (
     _resnet_fpn_extractor,
     _validate_trainable_layers,
@@ -27,6 +29,7 @@ class MaskRCNNResNet50FPNWeights(Weights):
         transforms=CocoEval,
         meta={
             "categories": _COCO_CATEGORIES,
+            "interpolation": InterpolationMode.BILINEAR,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn",
             "box_map": 37.9,
             "mask_map": 34.6,
diff --git a/torchvision/prototype/models/detection/retinanet.py b/torchvision/prototype/models/detection/retinanet.py
@@ -0,0 +1,77 @@
+import warnings
+from typing import Any, Optional
+
+from torchvision.transforms.functional import InterpolationMode
+
+from ....models.detection.retinanet import (
+    _resnet_fpn_extractor,
+    _validate_trainable_layers,
+    RetinaNet,
+    LastLevelP6P7,
+    misc_nn_ops,
+    overwrite_eps,
+)
+from ...transforms.presets import CocoEval
+from .._api import Weights, WeightEntry
+from .._meta import _COCO_CATEGORIES
+from ..resnet import ResNet50Weights, resnet50
+
+
+__all__ = [
+    "RetinaNet",
+    "RetinaNetResNet50FPNWeights",
+    "retinanet_resnet50_fpn",
+]
+
+
+class RetinaNetResNet50FPNWeights(Weights):
+    Coco_RefV1 = WeightEntry(
+        url="https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth",
+        transforms=CocoEval,
+        meta={
+            "categories": _COCO_CATEGORIES,
+            "interpolation": InterpolationMode.BILINEAR,
+            "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#retinanet",
+            "map": 36.4,
+        },
+    )
+
+
+def retinanet_resnet50_fpn(
+    weights: Optional[RetinaNetResNet50FPNWeights] = None,
+    weights_backbone: Optional[ResNet50Weights] = None,
+    progress: bool = True,
+    num_classes: int = 91,
+    trainable_backbone_layers: Optional[int] = None,
+    **kwargs: Any,
+) -> RetinaNet:
+    if "pretrained" in kwargs:
+        warnings.warn("The argument pretrained is deprecated, please use weights instead.")
+        weights = RetinaNetResNet50FPNWeights.Coco_RefV1 if kwargs.pop("pretrained") else None
+    weights = RetinaNetResNet50FPNWeights.verify(weights)
+    if "pretrained_backbone" in kwargs:
+        warnings.warn("The argument pretrained_backbone is deprecated, please use weights_backbone instead.")
+        weights_backbone = ResNet50Weights.ImageNet1K_RefV1 if kwargs.pop("pretrained_backbone") else None
+    weights_backbone = ResNet50Weights.verify(weights_backbone)
+
+    if weights is not None:
+        weights_backbone = None
+        num_classes = len(weights.meta["categories"])
+
+    trainable_backbone_layers = _validate_trainable_layers(
+        weights is not None or weights_backbone is not None, trainable_backbone_layers, 5, 3
+    )
+
+    backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=misc_nn_ops.FrozenBatchNorm2d)
+    # skip P2 because it generates too many anchors (according to their paper)
+    backbone = _resnet_fpn_extractor(
+        backbone, trainable_backbone_layers, returned_layers=[2, 3, 4], extra_blocks=LastLevelP6P7(256, 256)
+    )
+    model = RetinaNet(backbone, num_classes, **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.state_dict(progress=progress))
+        if weights == RetinaNetResNet50FPNWeights.Coco_RefV1:
+            overwrite_eps(model, 0.0)
+
+    return model
diff --git a/torchvision/prototype/models/segmentation/deeplabv3.py b/torchvision/prototype/models/segmentation/deeplabv3.py
@@ -2,6 +2,8 @@
 from functools import partial
 from typing import Any, Optional
 
+from torchvision.transforms.functional import InterpolationMode
+
 from ....models.segmentation.deeplabv3 import DeepLabV3, _deeplabv3_mobilenetv3, _deeplabv3_resnet
 from ...transforms.presets import VocEval
 from .._api import Weights, WeightEntry
@@ -22,7 +24,10 @@
 ]
 
 
-_common_meta = {"categories": _VOC_CATEGORIES}
+_common_meta = {
+    "categories": _VOC_CATEGORIES,
+    "interpolation": InterpolationMode.BILINEAR,
+}
 
 
 class DeepLabV3ResNet50Weights(Weights):
diff --git a/torchvision/prototype/models/segmentation/fcn.py b/torchvision/prototype/models/segmentation/fcn.py
@@ -2,6 +2,8 @@
 from functools import partial
 from typing import Any, Optional
 
+from torchvision.transforms.functional import InterpolationMode
+
 from ....models.segmentation.fcn import FCN, _fcn_resnet
 from ...transforms.presets import VocEval
 from .._api import Weights, WeightEntry
@@ -12,7 +14,10 @@
 __all__ = ["FCN", "FCNResNet50Weights", "FCNResNet101Weights", "fcn_resnet50", "fcn_resnet101"]
 
 
-_common_meta = {"categories": _VOC_CATEGORIES}
+_common_meta = {
+    "categories": _VOC_CATEGORIES,
+    "interpolation": InterpolationMode.BILINEAR,
+}
 
 
 class FCNResNet50Weights(Weights):
diff --git a/torchvision/prototype/models/segmentation/lraspp.py b/torchvision/prototype/models/segmentation/lraspp.py
@@ -2,6 +2,8 @@
 from functools import partial
 from typing import Any, Optional
 
+from torchvision.transforms.functional import InterpolationMode
+
 from ....models.segmentation.lraspp import LRASPP, _lraspp_mobilenetv3
 from ...transforms.presets import VocEval
 from .._api import Weights, WeightEntry
@@ -18,6 +20,7 @@ class LRASPPMobileNetV3LargeWeights(Weights):
         transforms=partial(VocEval, resize_size=520),
         meta={
             "categories": _VOC_CATEGORIES,
+            "interpolation": InterpolationMode.BILINEAR,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#lraspp_mobilenet_v3_large",
             "mIoU": 57.9,
             "acc": 91.2,
diff --git a/torchvision/prototype/models/vgg.py b/torchvision/prototype/models/vgg.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,3 @@`
`1`		`-import os`
`2`	`1`	`import os.path`
`3`	`2`	`import pickle`
`4`	`3`	`from typing import Any, Callable, Optional, Tuple`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`import io`
`2`		`-import os`
`3`	`2`	`import os.path`
`4`	`3`	`import pickle`
`5`	`4`	`import string`