From 68dc2a9034dbd76d21b30da572a61de98afbad46 Mon Sep 17 00:00:00 2001 From: zhiqwang Date: Mon, 14 Dec 2020 22:07:40 -0500 Subject: [PATCH 1/2] Replacing all torch.jit.annotations with typing --- test/test_ops.py | 2 +- torchvision/models/detection/_utils.py | 2 +- torchvision/models/detection/anchor_utils.py | 4 ++-- torchvision/models/detection/generalized_rcnn.py | 8 +++----- torchvision/models/detection/image_list.py | 2 +- torchvision/models/detection/retinanet.py | 11 +++++------ torchvision/models/detection/roi_heads.py | 8 ++++---- torchvision/models/detection/rpn.py | 2 +- torchvision/models/detection/transform.py | 4 ++-- torchvision/models/inception.py | 7 +++---- torchvision/models/quantization/googlenet.py | 1 - torchvision/models/quantization/inception.py | 1 - torchvision/ops/_box_convert.py | 2 -- torchvision/ops/_utils.py | 2 +- torchvision/ops/boxes.py | 2 +- torchvision/ops/deform_conv.py | 2 +- torchvision/ops/feature_pyramid_network.py | 3 +-- torchvision/ops/misc.py | 4 ++-- torchvision/ops/poolers.py | 9 +++------ torchvision/ops/ps_roi_align.py | 1 - torchvision/ops/roi_align.py | 2 +- torchvision/ops/roi_pool.py | 2 +- torchvision/transforms/autoaugment.py | 3 +-- torchvision/transforms/functional.py | 3 +-- torchvision/transforms/functional_tensor.py | 4 ++-- 25 files changed, 38 insertions(+), 53 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 1a83b4cf246..32d769a5654 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -7,9 +7,9 @@ import torch from torch import Tensor from torch.autograd import gradcheck -from torch.jit.annotations import Tuple from torch.nn.modules.utils import _pair from torchvision import ops +from typing import Tuple class OpTester(object): diff --git a/torchvision/models/detection/_utils.py b/torchvision/models/detection/_utils.py index 9834fb6b92a..93fd379a958 100644 --- a/torchvision/models/detection/_utils.py +++ b/torchvision/models/detection/_utils.py @@ -1,8 +1,8 @@ import math import torch -from torch.jit.annotations import List, Tuple from torch import Tensor +from typing import List, Tuple from torchvision.ops.misc import FrozenBatchNorm2d diff --git a/torchvision/models/detection/anchor_utils.py b/torchvision/models/detection/anchor_utils.py index defee6e5ae2..84ba9105e1c 100644 --- a/torchvision/models/detection/anchor_utils.py +++ b/torchvision/models/detection/anchor_utils.py @@ -2,7 +2,7 @@ import torch from torch import nn, Tensor -from torch.jit.annotations import List, Optional, Dict +from typing import List, Optional, Dict from .image_list import ImageList @@ -148,7 +148,7 @@ def forward(self, image_list: ImageList, feature_maps: List[Tensor]) -> List[Ten torch.tensor(image_size[1] // g[1], dtype=torch.int64, device=device)] for g in grid_sizes] self.set_cell_anchors(dtype, device) anchors_over_all_feature_maps = self.cached_grid_anchors(grid_sizes, strides) - anchors = torch.jit.annotate(List[List[torch.Tensor]], []) + anchors: List[List[torch.Tensor]] = [] for i in range(len(image_list.image_sizes)): anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps] anchors.append(anchors_in_image) diff --git a/torchvision/models/detection/generalized_rcnn.py b/torchvision/models/detection/generalized_rcnn.py index 47e7bbedfd4..e259920df11 100644 --- a/torchvision/models/detection/generalized_rcnn.py +++ b/torchvision/models/detection/generalized_rcnn.py @@ -4,12 +4,10 @@ """ from collections import OrderedDict -from typing import Union import torch -from torch import nn +from torch import nn, Tensor import warnings -from torch.jit.annotations import Tuple, List, Dict, Optional -from torch import Tensor +from typing import Tuple, List, Dict, Optional, Union class GeneralizedRCNN(nn.Module): @@ -71,7 +69,7 @@ def forward(self, images, targets=None): raise ValueError("Expected target boxes to be of type " "Tensor, got {:}.".format(type(boxes))) - original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], []) + original_image_sizes: List[Tuple[int, int]] = [] for img in images: val = img.shape[-2:] assert len(val) == 2 diff --git a/torchvision/models/detection/image_list.py b/torchvision/models/detection/image_list.py index 21e326573bf..e78f356c02a 100644 --- a/torchvision/models/detection/image_list.py +++ b/torchvision/models/detection/image_list.py @@ -1,7 +1,7 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. import torch -from torch.jit.annotations import List, Tuple from torch import Tensor +from typing import List, Tuple class ImageList(object): diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py index ee647ab3b0b..447a9156a34 100644 --- a/torchvision/models/detection/retinanet.py +++ b/torchvision/models/detection/retinanet.py @@ -3,9 +3,8 @@ import warnings import torch -import torch.nn as nn -from torch import Tensor -from torch.jit.annotations import Dict, List, Tuple, Optional +from torch import nn, Tensor +from typing import Dict, List, Tuple, Optional from ._utils import overwrite_eps from ..utils import load_state_dict_from_url @@ -402,7 +401,7 @@ def postprocess_detections(self, head_outputs, anchors, image_shapes): num_images = len(image_shapes) - detections = torch.jit.annotate(List[Dict[str, Tensor]], []) + detections: List[Dict[str, Tensor]] = [] for index in range(num_images): box_regression_per_image = [br[index] for br in box_regression] @@ -486,7 +485,7 @@ def forward(self, images, targets=None): "Tensor, got {:}.".format(type(boxes))) # get the original image sizes - original_image_sizes = torch.jit.annotate(List[Tuple[int, int]], []) + original_image_sizes: List[Tuple[int, int]] = [] for img in images: val = img.shape[-2:] assert len(val) == 2 @@ -524,7 +523,7 @@ def forward(self, images, targets=None): anchors = self.anchor_generator(images, features) losses = {} - detections = torch.jit.annotate(List[Dict[str, Tensor]], []) + detections: List[Dict[str, Tensor]] = [] if self.training: assert targets is not None diff --git a/torchvision/models/detection/roi_heads.py b/torchvision/models/detection/roi_heads.py index d67d5856f76..84584c7f77d 100644 --- a/torchvision/models/detection/roi_heads.py +++ b/torchvision/models/detection/roi_heads.py @@ -10,7 +10,7 @@ from . import _utils as det_utils -from torch.jit.annotations import Optional, List, Dict, Tuple +from typing import Optional, List, Dict, Tuple def fastrcnn_loss(class_logits, box_regression, labels, regression_targets): @@ -379,7 +379,7 @@ def expand_masks(mask, padding): scale = expand_masks_tracing_scale(M, padding) else: scale = float(M + 2 * padding) / M - padded_mask = torch.nn.functional.pad(mask, (padding,) * 4) + padded_mask = F.pad(mask, (padding,) * 4) return padded_mask, scale @@ -482,7 +482,7 @@ def paste_masks_in_image(masks, boxes, img_shape, padding=1): return ret -class RoIHeads(torch.nn.Module): +class RoIHeads(nn.Module): __annotations__ = { 'box_coder': det_utils.BoxCoder, 'proposal_matcher': det_utils.Matcher, @@ -753,7 +753,7 @@ def forward(self, box_features = self.box_head(box_features) class_logits, box_regression = self.box_predictor(box_features) - result = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) + result: List[Dict[str, torch.Tensor]] = [] losses = {} if self.training: assert labels is not None and regression_targets is not None diff --git a/torchvision/models/detection/rpn.py b/torchvision/models/detection/rpn.py index e4c4478eb54..72488853d72 100644 --- a/torchvision/models/detection/rpn.py +++ b/torchvision/models/detection/rpn.py @@ -9,7 +9,7 @@ from . import _utils as det_utils from .image_list import ImageList -from torch.jit.annotations import List, Optional, Dict, Tuple +from typing import List, Optional, Dict, Tuple # Import AnchorGenerator to keep compatibility. from .anchor_utils import AnchorGenerator diff --git a/torchvision/models/detection/transform.py b/torchvision/models/detection/transform.py index 9059c184949..8b91d536621 100644 --- a/torchvision/models/detection/transform.py +++ b/torchvision/models/detection/transform.py @@ -4,7 +4,7 @@ from torch import nn, Tensor from torch.nn import functional as F import torchvision -from torch.jit.annotations import List, Tuple, Dict, Optional +from typing import List, Tuple, Dict, Optional from .image_list import ImageList from .roi_heads import paste_masks_in_image @@ -109,7 +109,7 @@ def forward(self, image_sizes = [img.shape[-2:] for img in images] images = self.batch_images(images) - image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], []) + image_sizes_list: List[Tuple[int, int]] = [] for image_size in image_sizes: assert len(image_size) == 2 image_sizes_list.append((image_size[0], image_size[1])) diff --git a/torchvision/models/inception.py b/torchvision/models/inception.py index f2156018327..4abad7c02e7 100644 --- a/torchvision/models/inception.py +++ b/torchvision/models/inception.py @@ -1,9 +1,8 @@ from collections import namedtuple import warnings import torch -import torch.nn as nn +from torch import nn, Tensor import torch.nn.functional as F -from torch import Tensor from .utils import load_state_dict_from_url from typing import Callable, Any, Optional, Tuple, List @@ -17,7 +16,7 @@ } InceptionOutputs = namedtuple('InceptionOutputs', ['logits', 'aux_logits']) -InceptionOutputs.__annotations__ = {'logits': torch.Tensor, 'aux_logits': Optional[torch.Tensor]} +InceptionOutputs.__annotations__ = {'logits': Tensor, 'aux_logits': Optional[Tensor]} # Script annotations failed with _GoogleNetOutputs = namedtuple ... # _InceptionOutputs set here for backwards compat @@ -171,7 +170,7 @@ def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor]]: # N x 768 x 17 x 17 x = self.Mixed_6e(x) # N x 768 x 17 x 17 - aux = torch.jit.annotate(Optional[Tensor], None) + aux: Optional[Tensor] = None if self.AuxLogits is not None: if self.training: aux = self.AuxLogits(x) diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py index d01534bc70c..54ee6db0eae 100644 --- a/torchvision/models/quantization/googlenet.py +++ b/torchvision/models/quantization/googlenet.py @@ -2,7 +2,6 @@ import torch import torch.nn as nn from torch.nn import functional as F -from torch.jit.annotations import Optional from torchvision.models.utils import load_state_dict_from_url from torchvision.models.googlenet import ( diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py index cdc72ce4851..ae588e38256 100644 --- a/torchvision/models/quantization/inception.py +++ b/torchvision/models/quantization/inception.py @@ -6,7 +6,6 @@ import torch.nn.functional as F from torchvision.models import inception as inception_module from torchvision.models.inception import InceptionOutputs -from torch.jit.annotations import Optional from torchvision.models.utils import load_state_dict_from_url from .utils import _replace_relu, quantize_model diff --git a/torchvision/ops/_box_convert.py b/torchvision/ops/_box_convert.py index 121ec0427bd..0d90cc36039 100644 --- a/torchvision/ops/_box_convert.py +++ b/torchvision/ops/_box_convert.py @@ -1,7 +1,5 @@ import torch -from torch.jit.annotations import Tuple from torch import Tensor -import torchvision def _box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor: diff --git a/torchvision/ops/_utils.py b/torchvision/ops/_utils.py index cc964fd0e33..bb6287ad616 100644 --- a/torchvision/ops/_utils.py +++ b/torchvision/ops/_utils.py @@ -1,6 +1,6 @@ import torch from torch import Tensor -from torch.jit.annotations import List +from typing import List def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor: diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index 8c6e2ec6da4..b7c0761fc14 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -1,6 +1,6 @@ import torch -from torch.jit.annotations import Tuple from torch import Tensor +from typing import Tuple from ._box_convert import _box_cxcywh_to_xyxy, _box_xyxy_to_cxcywh, _box_xywh_to_xyxy, _box_xyxy_to_xywh import torchvision from torchvision.extension import _assert_has_ops diff --git a/torchvision/ops/deform_conv.py b/torchvision/ops/deform_conv.py index 54d1bfdc152..0f0c5fe1e42 100644 --- a/torchvision/ops/deform_conv.py +++ b/torchvision/ops/deform_conv.py @@ -5,7 +5,7 @@ from torch.nn import init from torch.nn.parameter import Parameter from torch.nn.modules.utils import _pair -from torch.jit.annotations import Optional, Tuple +from typing import Optional, Tuple from torchvision.extension import _assert_has_ops diff --git a/torchvision/ops/feature_pyramid_network.py b/torchvision/ops/feature_pyramid_network.py index b68e4b775d6..b0c301ec474 100644 --- a/torchvision/ops/feature_pyramid_network.py +++ b/torchvision/ops/feature_pyramid_network.py @@ -1,10 +1,9 @@ from collections import OrderedDict -import torch import torch.nn.functional as F from torch import nn, Tensor -from torch.jit.annotations import Tuple, List, Dict, Optional +from typing import Tuple, List, Dict, Optional class ExtraFPNBlock(nn.Module): diff --git a/torchvision/ops/misc.py b/torchvision/ops/misc.py index 3e9f13c9daf..7e43caa78d6 100644 --- a/torchvision/ops/misc.py +++ b/torchvision/ops/misc.py @@ -10,8 +10,8 @@ import warnings import torch -from torch import Tensor, Size -from torch.jit.annotations import List, Optional, Tuple +from torch import Tensor +from typing import List, Optional class Conv2d(torch.nn.Conv2d): diff --git a/torchvision/ops/poolers.py b/torchvision/ops/poolers.py index 02dbf3904bb..a5e74fb99f5 100644 --- a/torchvision/ops/poolers.py +++ b/torchvision/ops/poolers.py @@ -1,15 +1,12 @@ # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. -from typing import Union - import torch -import torch.nn.functional as F from torch import nn, Tensor +import torchvision from torchvision.ops import roi_align from torchvision.ops.boxes import box_area -from torch.jit.annotations import Optional, List, Dict, Tuple -import torchvision +from typing import Optional, List, Dict, Tuple, Union # copying result_idx_in_level to a specific index in result[] @@ -149,7 +146,7 @@ def convert_to_roi_format(self, boxes: List[Tensor]) -> Tensor: def infer_scale(self, feature: Tensor, original_size: List[int]) -> float: # assumption: the scale is of the form 2 ** (-k), with k integer size = feature.shape[-2:] - possible_scales = torch.jit.annotate(List[float], []) + possible_scales: List[float] = [] for s1, s2 in zip(size, original_size): approx_scale = float(s1) / float(s2) scale = 2 ** float(torch.tensor(approx_scale).log2().round()) diff --git a/torchvision/ops/ps_roi_align.py b/torchvision/ops/ps_roi_align.py index 8e0180f62ce..e8a2ff7aaeb 100644 --- a/torchvision/ops/ps_roi_align.py +++ b/torchvision/ops/ps_roi_align.py @@ -2,7 +2,6 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, Tuple from torchvision.extension import _assert_has_ops from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py index e32416c3b50..a92e1df7fd7 100644 --- a/torchvision/ops/roi_align.py +++ b/torchvision/ops/roi_align.py @@ -2,7 +2,7 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2 +from torch.jit.annotations import BroadcastingList2 from torchvision.extension import _assert_has_ops from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py index f7d778e6ba4..b80111ca4a5 100644 --- a/torchvision/ops/roi_pool.py +++ b/torchvision/ops/roi_pool.py @@ -2,7 +2,7 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2 +from torch.jit.annotations import BroadcastingList2 from torchvision.extension import _assert_has_ops from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape diff --git a/torchvision/transforms/autoaugment.py b/torchvision/transforms/autoaugment.py index 26847521998..ab6d0a69e5c 100644 --- a/torchvision/transforms/autoaugment.py +++ b/torchvision/transforms/autoaugment.py @@ -3,8 +3,7 @@ from enum import Enum from torch import Tensor -from torch.jit.annotations import List, Tuple -from typing import Optional +from typing import List, Tuple, Optional from . import functional as F, InterpolationMode diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index ec7e511989a..2fb74269195 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -2,14 +2,13 @@ import numbers import warnings from enum import Enum -from typing import Any, Optional import numpy as np from PIL import Image import torch from torch import Tensor -from torch.jit.annotations import List, Tuple +from typing import List, Tuple, Any, Optional try: import accimage diff --git a/torchvision/transforms/functional_tensor.py b/torchvision/transforms/functional_tensor.py index a72cc41f5cd..7203e79d06b 100644 --- a/torchvision/transforms/functional_tensor.py +++ b/torchvision/transforms/functional_tensor.py @@ -1,10 +1,10 @@ import warnings -from typing import Optional, Tuple import torch from torch import Tensor from torch.nn.functional import grid_sample, conv2d, interpolate, pad as torch_pad -from torch.jit.annotations import List, BroadcastingList2 +from torch.jit.annotations import BroadcastingList2 +from typing import Optional, Tuple, List def _is_tensor_a_torch_image(x: Tensor) -> bool: From 5b58149bbafd4b574a9e62df14fd2b73faef7976 Mon Sep 17 00:00:00 2001 From: zhiqwang Date: Mon, 14 Dec 2020 22:21:23 -0500 Subject: [PATCH 2/2] Replacing remaining typing --- torchvision/models/googlenet.py | 4 ++-- torchvision/ops/ps_roi_pool.py | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/torchvision/models/googlenet.py b/torchvision/models/googlenet.py index 274ac74bc04..cef48dea76a 100644 --- a/torchvision/models/googlenet.py +++ b/torchvision/models/googlenet.py @@ -162,7 +162,7 @@ def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor # N x 480 x 14 x 14 x = self.inception4a(x) # N x 512 x 14 x 14 - aux1 = torch.jit.annotate(Optional[Tensor], None) + aux1: Optional[Tensor] = None if self.aux1 is not None: if self.training: aux1 = self.aux1(x) @@ -173,7 +173,7 @@ def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor # N x 512 x 14 x 14 x = self.inception4d(x) # N x 528 x 14 x 14 - aux2 = torch.jit.annotate(Optional[Tensor], None) + aux2: Optional[Tensor] = None if self.aux2 is not None: if self.training: aux2 = self.aux2(x) diff --git a/torchvision/ops/ps_roi_pool.py b/torchvision/ops/ps_roi_pool.py index bf08d84e78f..c44cf937b57 100644 --- a/torchvision/ops/ps_roi_pool.py +++ b/torchvision/ops/ps_roi_pool.py @@ -2,7 +2,6 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, Tuple from torchvision.extension import _assert_has_ops from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape