From 876ddd4f1da5c5f32409f26e57d754b1dc579d0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vuka=C5=A1in=20Manojlovi=C4=87?= Date: Thu, 18 Jun 2020 15:00:40 +0200 Subject: [PATCH 1/4] Add type annotations for torchvision.ops --- torchvision/ops/_utils.py | 8 ++-- torchvision/ops/boxes.py | 17 +++++--- torchvision/ops/deform_conv.py | 35 +++++++++++----- torchvision/ops/feature_pyramid_network.py | 38 ++++++++++++----- torchvision/ops/misc.py | 25 +++++++++--- torchvision/ops/new_empty_tensor.py | 2 +- torchvision/ops/poolers.py | 47 +++++++++++++++++----- torchvision/ops/ps_roi_align.py | 21 +++++++--- torchvision/ops/ps_roi_pool.py | 15 ++++--- torchvision/ops/roi_align.py | 23 ++++++++--- torchvision/ops/roi_pool.py | 15 ++++--- 11 files changed, 182 insertions(+), 64 deletions(-) diff --git a/torchvision/ops/_utils.py b/torchvision/ops/_utils.py index f514664042b..37751594a06 100644 --- a/torchvision/ops/_utils.py +++ b/torchvision/ops/_utils.py @@ -1,9 +1,9 @@ import torch from torch import Tensor -from torch.jit.annotations import List +from torch.jit.annotations import List, Tuple, Union -def _cat(tensors, dim=0): +def _cat(tensors: Union[List[Tensor], Tuple[Tensor, ...]], dim: int = 0) -> Tensor: # type: (List[Tensor], int) -> Tensor """ Efficient version of torch.cat that avoids a copy if there is only a single element in a list @@ -15,7 +15,7 @@ def _cat(tensors, dim=0): return torch.cat(tensors, dim) -def convert_boxes_to_roi_format(boxes): +def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor: # type: (List[Tensor]) -> Tensor concat_boxes = _cat([b for b in boxes], dim=0) temp = [] @@ -26,7 +26,7 @@ def convert_boxes_to_roi_format(boxes): return rois -def check_roi_boxes_shape(boxes): +def check_roi_boxes_shape(boxes: Union[List[Tensor], Tuple[Tensor, ...], Tensor]): if isinstance(boxes, (list, tuple)): for _tensor in boxes: assert _tensor.size(1) == 4, \ diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index e7442f57352..9039605f4fa 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -5,7 +5,7 @@ @torch.jit.script -def nms(boxes, scores, iou_threshold): +def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: # type: (Tensor, Tensor, float) -> Tensor """ Performs non-maximum suppression (NMS) on the boxes according @@ -42,7 +42,12 @@ def nms(boxes, scores, iou_threshold): @torch.jit.script -def batched_nms(boxes, scores, idxs, iou_threshold): +def batched_nms( + boxes: Tensor, + scores: Tensor, + idxs: Tensor, + iou_threshold: float, +) -> Tensor: # type: (Tensor, Tensor, Tensor, float) -> Tensor """ Performs non-maximum suppression in a batched fashion. @@ -84,7 +89,7 @@ def batched_nms(boxes, scores, idxs, iou_threshold): return keep -def remove_small_boxes(boxes, min_size): +def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor: # type: (Tensor, float) -> Tensor """ Remove boxes which contains at least one side smaller than min_size. @@ -103,7 +108,7 @@ def remove_small_boxes(boxes, min_size): return keep -def clip_boxes_to_image(boxes, size): +def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor: # type: (Tensor, Tuple[int, int]) -> Tensor """ Clip boxes so that they lie inside an image of size `size`. @@ -133,7 +138,7 @@ def clip_boxes_to_image(boxes, size): return clipped_boxes.reshape(boxes.shape) -def box_area(boxes): +def box_area(boxes: Tensor) -> Tensor: """ Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates. @@ -150,7 +155,7 @@ def box_area(boxes): # implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py # with slight modifications -def box_iou(boxes1, boxes2): +def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor: """ Return intersection-over-union (Jaccard index) of boxes. diff --git a/torchvision/ops/deform_conv.py b/torchvision/ops/deform_conv.py index c948b164196..409d5ba413f 100644 --- a/torchvision/ops/deform_conv.py +++ b/torchvision/ops/deform_conv.py @@ -5,10 +5,18 @@ 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 - - -def deform_conv2d(input, offset, weight, bias=None, stride=(1, 1), padding=(0, 0), dilation=(1, 1)): +from torch.jit.annotations import Optional, Tuple, Union + + +def deform_conv2d( + input: Tensor, + offset: Tensor, + weight: Tensor, + bias: Optional[Tensor] = None, + stride: Union[int, Tuple[int, int]] = (1, 1), + padding: Union[int, Tuple[int, int]] = (0, 0), + dilation: Union[int, Tuple[int, int]] = (1, 1), +) -> Tensor: # type: (Tensor, Tensor, Tensor, Optional[Tensor], Tuple[int, int], Tuple[int, int], Tuple[int, int]) -> Tensor """ Performs Deformable Convolution, described in Deformable Convolutional Networks @@ -80,8 +88,17 @@ class DeformConv2d(nn.Module): """ See deform_conv2d """ - def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, - dilation=1, groups=1, bias=True): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]] = 1, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, + groups: int = 1, + bias: bool = True, + ): super(DeformConv2d, self).__init__() if in_channels % groups != 0: @@ -107,14 +124,14 @@ def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, self.reset_parameters() - def reset_parameters(self): + def reset_parameters(self) -> None: init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) init.uniform_(self.bias, -bound, bound) - def forward(self, input, offset): + def forward(self, input: Tensor, offset: Tensor) -> Tensor: """ Arguments: input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor @@ -125,7 +142,7 @@ def forward(self, input, offset): return deform_conv2d(input, offset, self.weight, self.bias, stride=self.stride, padding=self.padding, dilation=self.dilation) - def __repr__(self): + def __repr__(self) -> str: s = self.__class__.__name__ + '(' s += '{in_channels}' s += ', {out_channels}' diff --git a/torchvision/ops/feature_pyramid_network.py b/torchvision/ops/feature_pyramid_network.py index a2d8c409490..a55431f4efd 100644 --- a/torchvision/ops/feature_pyramid_network.py +++ b/torchvision/ops/feature_pyramid_network.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from torch import nn, Tensor -from torch.jit.annotations import Tuple, List, Dict +from torch.jit.annotations import Tuple, List, Dict, Optional class FeaturePyramidNetwork(nn.Module): @@ -44,7 +44,12 @@ class FeaturePyramidNetwork(nn.Module): >>> ('feat3', torch.Size([1, 5, 8, 8]))] """ - def __init__(self, in_channels_list, out_channels, extra_blocks=None): + def __init__( + self, + in_channels_list: List[int], + out_channels: int, + extra_blocks: Optional[ExtraFPNBlock] = None, + ): super(FeaturePyramidNetwork, self).__init__() self.inner_blocks = nn.ModuleList() self.layer_blocks = nn.ModuleList() @@ -66,7 +71,7 @@ def __init__(self, in_channels_list, out_channels, extra_blocks=None): assert isinstance(extra_blocks, ExtraFPNBlock) self.extra_blocks = extra_blocks - def get_result_from_inner_blocks(self, x, idx): + def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: # type: (Tensor, int) -> Tensor """ This is equivalent to self.inner_blocks[idx](x), @@ -85,7 +90,7 @@ def get_result_from_inner_blocks(self, x, idx): i += 1 return out - def get_result_from_layer_blocks(self, x, idx): + def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: # type: (Tensor, int) -> Tensor """ This is equivalent to self.layer_blocks[idx](x), @@ -104,7 +109,7 @@ def get_result_from_layer_blocks(self, x, idx): i += 1 return out - def forward(self, x): + def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: # type: (Dict[str, Tensor]) -> Dict[str, Tensor] """ Computes the FPN for a set of feature maps. @@ -155,7 +160,12 @@ class ExtraFPNBlock(nn.Module): of the FPN names (List[str]): the extended set of names for the results """ - def forward(self, results, x, names): + def forward( + self, + results: List[Tensor], + x: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: pass @@ -163,7 +173,12 @@ class LastLevelMaxPool(ExtraFPNBlock): """ Applies a max_pool2d on top of the last feature map """ - def forward(self, x, y, names): + def forward( + self, + x: List[Tensor], + y: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: # type: (List[Tensor], List[Tensor], List[str]) -> Tuple[List[Tensor], List[str]] names.append("pool") x.append(F.max_pool2d(x[-1], 1, 2, 0)) @@ -174,7 +189,7 @@ class LastLevelP6P7(ExtraFPNBlock): """ This module is used in RetinaNet to generate extra layers, P6 and P7. """ - def __init__(self, in_channels, out_channels): + def __init__(self, in_channels: int, out_channels: int): super(LastLevelP6P7, self).__init__() self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1) self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1) @@ -183,7 +198,12 @@ def __init__(self, in_channels, out_channels): nn.init.constant_(module.bias, 0) self.use_P5 = in_channels == out_channels - def forward(self, p, c, names): + def forward( + self, + p: List[Tensor], + c: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: p5, c5 = p[-1], c[-1] x = p5 if self.use_P5 else c5 p6 = self.p6(x) diff --git a/torchvision/ops/misc.py b/torchvision/ops/misc.py index 61fab3edd7a..6931e325567 100644 --- a/torchvision/ops/misc.py +++ b/torchvision/ops/misc.py @@ -10,6 +10,8 @@ import warnings import torch +from torch import Tensor, Size +from torch.jit.annotations import List, Optional, Tuple, Union class Conv2d(torch.nn.Conv2d): @@ -46,7 +48,12 @@ class FrozenBatchNorm2d(torch.nn.Module): are fixed """ - def __init__(self, num_features, eps=0., n=None): + def __init__( + self, + num_features: Union[Size, List[int], Tuple[int, ...]], + eps: float = 0., + n: Optional[Size, Union[List[int], Tuple[int, ...]]] = None, + ): # n=None for backward-compatibility if n is not None: warnings.warn("`n` argument is deprecated and has been renamed `num_features`", @@ -59,8 +66,16 @@ def __init__(self, num_features, eps=0., n=None): self.register_buffer("running_mean", torch.zeros(num_features)) self.register_buffer("running_var", torch.ones(num_features)) - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs): + def _load_from_state_dict( + self, + state_dict: dict, + prefix: str, + local_metadata: dict, + strict: bool, + missing_keys: List[str], + unexpected_keys: List[str], + error_msgs: List[str], + ): num_batches_tracked_key = prefix + 'num_batches_tracked' if num_batches_tracked_key in state_dict: del state_dict[num_batches_tracked_key] @@ -69,7 +84,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: # move reshapes to the beginning # to make it fuser-friendly w = self.weight.reshape(1, -1, 1, 1) @@ -80,5 +95,5 @@ def forward(self, x): bias = b - rm * scale return x * scale + bias - def __repr__(self): + def __repr__(self) -> str: return f"{self.__class__.__name__}({self.weight.shape[0]})" diff --git a/torchvision/ops/new_empty_tensor.py b/torchvision/ops/new_empty_tensor.py index 74455a98c4f..fb3f2f4ee8c 100644 --- a/torchvision/ops/new_empty_tensor.py +++ b/torchvision/ops/new_empty_tensor.py @@ -3,7 +3,7 @@ from torch import Tensor -def _new_empty_tensor(x, shape): +def _new_empty_tensor(x: Tensor, shape: List[int]) -> Tensor: # type: (Tensor, List[int]) -> Tensor """ Arguments: diff --git a/torchvision/ops/poolers.py b/torchvision/ops/poolers.py index 06bbc86a93c..7d3f2be7994 100644 --- a/torchvision/ops/poolers.py +++ b/torchvision/ops/poolers.py @@ -6,7 +6,7 @@ from torchvision.ops import roi_align from torchvision.ops.boxes import box_area -from torch.jit.annotations import Optional, List, Dict, Tuple +from torch.jit.annotations import Optional, List, Dict, Tuple, Union import torchvision @@ -15,7 +15,7 @@ # _onnx_merge_levels() is an implementation supported by ONNX # that merges the levels to the right indices @torch.jit.unused -def _onnx_merge_levels(levels, unmerged_results): +def _onnx_merge_levels(levels: Tensor, unmerged_results: List[Tensor]) -> Tensor: # type: (Tensor, List[Tensor]) -> Tensor first_result = unmerged_results[0] dtype, device = first_result.dtype, first_result.device @@ -33,7 +33,13 @@ def _onnx_merge_levels(levels, unmerged_results): # TODO: (eellison) T54974082 https://github.com/pytorch/pytorch/issues/26744/pytorch/issues/26744 -def initLevelMapper(k_min, k_max, canonical_scale=224, canonical_level=4, eps=1e-6): +def initLevelMapper( + k_min: int, + k_max: int, + canonical_scale: int = 224, + canonical_level: int = 4, + eps: float = 1e-6, +): # type: (int, int, int, int, float) -> LevelMapper return LevelMapper(k_min, k_max, canonical_scale, canonical_level, eps) @@ -50,7 +56,14 @@ class LevelMapper(object): eps (float) """ - def __init__(self, k_min, k_max, canonical_scale=224, canonical_level=4, eps=1e-6): + def __init__( + self, + k_min: int, + k_max: int, + canonical_scale: int = 224, + canonical_level: int = 4, + eps: float = 1e-6, + ): # type: (int, int, int, int, float) -> None self.k_min = k_min self.k_max = k_max @@ -58,7 +71,7 @@ def __init__(self, k_min, k_max, canonical_scale=224, canonical_level=4, eps=1e- self.lvl0 = canonical_level self.eps = eps - def __call__(self, boxlists): + def __call__(self, boxlists: List[Tensor]) -> Tensor: # type: (List[Tensor]) -> Tensor """ Arguments: @@ -107,7 +120,12 @@ class MultiScaleRoIAlign(nn.Module): 'map_levels': Optional[LevelMapper] } - def __init__(self, featmap_names, output_size, sampling_ratio): + def __init__( + self, + featmap_names: List[str], + output_size: Union[List[Tuple[int, int]], List[int]], + sampling_ratio: int, + ): super(MultiScaleRoIAlign, self).__init__() if isinstance(output_size, int): output_size = (output_size, output_size) @@ -117,7 +135,7 @@ def __init__(self, featmap_names, output_size, sampling_ratio): self.scales = None self.map_levels = None - def convert_to_roi_format(self, boxes): + def convert_to_roi_format(self, boxes: List[Tensor]) -> Tensor: # type: (List[Tensor]) -> Tensor concat_boxes = torch.cat(boxes, dim=0) device, dtype = concat_boxes.device, concat_boxes.dtype @@ -131,7 +149,7 @@ def convert_to_roi_format(self, boxes): rois = torch.cat([ids, concat_boxes], dim=1) return rois - def infer_scale(self, feature, original_size): + def infer_scale(self, feature: Tensor, original_size: List[int]) -> float: # type: (Tensor, List[int]) -> float # assumption: the scale is of the form 2 ** (-k), with k integer size = feature.shape[-2:] @@ -143,7 +161,11 @@ def infer_scale(self, feature, original_size): assert possible_scales[0] == possible_scales[1] return possible_scales[0] - def setup_scales(self, features, image_shapes): + def setup_scales( + self, + features: List[Tensor], + image_shapes: List[Tuple[int, int]], + ) -> None: # type: (List[Tensor], List[Tuple[int, int]]) -> None assert len(image_shapes) != 0 max_x = 0 @@ -161,7 +183,12 @@ def setup_scales(self, features, image_shapes): self.scales = scales self.map_levels = initLevelMapper(int(lvl_min), int(lvl_max)) - def forward(self, x, boxes, image_shapes): + def forward( + self, + x: Dict[str, Tensor], + boxes: List[Tensor], + image_shapes: List[Tuple[int, int]], + ) -> Tensor: # type: (Dict[str, Tensor], List[Tensor], List[Tuple[int, int]]) -> Tensor """ Arguments: diff --git a/torchvision/ops/ps_roi_align.py b/torchvision/ops/ps_roi_align.py index c0c761b72cc..26ba5bde88c 100644 --- a/torchvision/ops/ps_roi_align.py +++ b/torchvision/ops/ps_roi_align.py @@ -2,12 +2,18 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List +from torch.jit.annotations import List, Tuple, Union from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape -def ps_roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1): +def ps_roi_align( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: Union[int, Tuple[int, int]], + spatial_scale: float = 1.0, + sampling_ratio: int = -1, +) -> Tensor: # type: (Tensor, Tensor, int, float, int) -> Tensor """ Performs Position-Sensitive Region of Interest (RoI) Align operator @@ -49,17 +55,22 @@ class PSRoIAlign(nn.Module): """ See ps_roi_align """ - def __init__(self, output_size, spatial_scale, sampling_ratio): + def __init__( + self, + output_size: Union[int, Tuple[int, int]], + spatial_scale: float, + sampling_ratio: int, + ): super(PSRoIAlign, self).__init__() self.output_size = output_size self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio - def forward(self, input, rois): + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: return ps_roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio) - def __repr__(self): + def __repr__(self) -> str: tmpstr = self.__class__.__name__ + '(' tmpstr += 'output_size=' + str(self.output_size) tmpstr += ', spatial_scale=' + str(self.spatial_scale) diff --git a/torchvision/ops/ps_roi_pool.py b/torchvision/ops/ps_roi_pool.py index 710f2cb0195..e4b4cf105a5 100644 --- a/torchvision/ops/ps_roi_pool.py +++ b/torchvision/ops/ps_roi_pool.py @@ -2,12 +2,17 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List +from torch.jit.annotations import List, Tuple, Union from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape -def ps_roi_pool(input, boxes, output_size, spatial_scale=1.0): +def ps_roi_pool( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: Union[int, Tuple[int, int]], + spatial_scale: float = 1.0, +) -> Tensor: # type: (Tensor, Tensor, int, float) -> Tensor """ Performs Position-Sensitive Region of Interest (RoI) Pool operator @@ -43,15 +48,15 @@ class PSRoIPool(nn.Module): """ See ps_roi_pool """ - def __init__(self, output_size, spatial_scale): + def __init__(self, output_size: Union[int, Tuple[int, int]], spatial_scale: float): super(PSRoIPool, self).__init__() self.output_size = output_size self.spatial_scale = spatial_scale - def forward(self, input, rois): + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: return ps_roi_pool(input, rois, self.output_size, self.spatial_scale) - def __repr__(self): + def __repr__(self) -> str: tmpstr = self.__class__.__name__ + '(' tmpstr += 'output_size=' + str(self.output_size) tmpstr += ', spatial_scale=' + str(self.spatial_scale) diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py index 14224d8a83e..426e3c37e4d 100644 --- a/torchvision/ops/roi_align.py +++ b/torchvision/ops/roi_align.py @@ -2,12 +2,19 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2 +from torch.jit.annotations import List, BroadcastingList2, Union from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape -def roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1, aligned=False): +def roi_align( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: BroadcastingList2[int], + spatial_scale: float = 1.0, + sampling_ratio: int = -1, + aligned: bool = False, +) -> Tensor: # type: (Tensor, Tensor, BroadcastingList2[int], float, int, bool) -> Tensor """ Performs Region of Interest (RoI) Align operator described in Mask R-CNN @@ -49,17 +56,23 @@ class RoIAlign(nn.Module): """ See roi_align """ - def __init__(self, output_size, spatial_scale, sampling_ratio, aligned=False): + def __init__( + self, + output_size: BroadcastingList2[int], + spatial_scale: float, + sampling_ratio: int, + aligned: bool = False, + ): super(RoIAlign, self).__init__() self.output_size = output_size self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio self.aligned = aligned - def forward(self, input, rois): + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned) - def __repr__(self): + def __repr__(self) -> str: tmpstr = self.__class__.__name__ + '(' tmpstr += 'output_size=' + str(self.output_size) tmpstr += ', spatial_scale=' + str(self.spatial_scale) diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py index 10232f16b4a..ec79713cee2 100644 --- a/torchvision/ops/roi_pool.py +++ b/torchvision/ops/roi_pool.py @@ -2,12 +2,17 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2 +from torch.jit.annotations import List, BroadcastingList2, Union from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape -def roi_pool(input, boxes, output_size, spatial_scale=1.0): +def roi_pool( + input: Tensor, + boxes: Union[Tensor, List[Tensor]], + output_size: BroadcastingList2[int], + spatial_scale: float = 1.0, +) -> Tensor: # type: (Tensor, Tensor, BroadcastingList2[int], float) -> Tensor """ Performs Region of Interest (RoI) Pool operator described in Fast R-CNN @@ -41,15 +46,15 @@ class RoIPool(nn.Module): """ See roi_pool """ - def __init__(self, output_size, spatial_scale): + def __init__(self, output_size: BroadcastingList2[int], spatial_scale: float): super(RoIPool, self).__init__() self.output_size = output_size self.spatial_scale = spatial_scale - def forward(self, input, rois): + def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: return roi_pool(input, rois, self.output_size, self.spatial_scale) - def __repr__(self): + def __repr__(self) -> str: tmpstr = self.__class__.__name__ + '(' tmpstr += 'output_size=' + str(self.output_size) tmpstr += ', spatial_scale=' + str(self.spatial_scale) From b4c5f251f12968fbc3467c6929b42cfb84e56ad8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vuka=C5=A1in=20Manojlovi=C4=87?= Date: Tue, 30 Jun 2020 11:08:58 +0200 Subject: [PATCH 2/4] Fix type annotations for torchvision.ops --- torchvision/ops/_utils.py | 8 +++----- torchvision/ops/boxes.py | 4 ---- torchvision/ops/deform_conv.py | 17 ++++++++--------- torchvision/ops/feature_pyramid_network.py | 4 ---- torchvision/ops/misc.py | 6 +++--- torchvision/ops/new_empty_tensor.py | 1 - torchvision/ops/poolers.py | 14 +++----------- torchvision/ops/ps_roi_align.py | 11 +++++------ torchvision/ops/ps_roi_pool.py | 11 +++++------ torchvision/ops/roi_align.py | 7 +++---- torchvision/ops/roi_pool.py | 7 +++---- 11 files changed, 33 insertions(+), 57 deletions(-) diff --git a/torchvision/ops/_utils.py b/torchvision/ops/_utils.py index 37751594a06..6c9a040ecd0 100644 --- a/torchvision/ops/_utils.py +++ b/torchvision/ops/_utils.py @@ -1,10 +1,9 @@ import torch from torch import Tensor -from torch.jit.annotations import List, Tuple, Union +from torch.jit.annotations import List, Tuple -def _cat(tensors: Union[List[Tensor], Tuple[Tensor, ...]], dim: int = 0) -> Tensor: - # type: (List[Tensor], int) -> Tensor +def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor: """ Efficient version of torch.cat that avoids a copy if there is only a single element in a list """ @@ -16,7 +15,6 @@ def _cat(tensors: Union[List[Tensor], Tuple[Tensor, ...]], dim: int = 0) -> Tens def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor: - # type: (List[Tensor]) -> Tensor concat_boxes = _cat([b for b in boxes], dim=0) temp = [] for i, b in enumerate(boxes): @@ -26,7 +24,7 @@ def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor: return rois -def check_roi_boxes_shape(boxes: Union[List[Tensor], Tuple[Tensor, ...], Tensor]): +def check_roi_boxes_shape(boxes: Tensor): if isinstance(boxes, (list, tuple)): for _tensor in boxes: assert _tensor.size(1) == 4, \ diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index 9039605f4fa..5e43ca78a7f 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -6,7 +6,6 @@ @torch.jit.script def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: - # type: (Tensor, Tensor, float) -> Tensor """ Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU). @@ -48,7 +47,6 @@ def batched_nms( idxs: Tensor, iou_threshold: float, ) -> Tensor: - # type: (Tensor, Tensor, Tensor, float) -> Tensor """ Performs non-maximum suppression in a batched fashion. @@ -90,7 +88,6 @@ def batched_nms( def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor: - # type: (Tensor, float) -> Tensor """ Remove boxes which contains at least one side smaller than min_size. @@ -109,7 +106,6 @@ def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor: def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor: - # type: (Tensor, Tuple[int, int]) -> Tensor """ Clip boxes so that they lie inside an image of size `size`. diff --git a/torchvision/ops/deform_conv.py b/torchvision/ops/deform_conv.py index 409d5ba413f..aa5e42c4a6e 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, Union +from torch.jit.annotations import Optional, Tuple def deform_conv2d( @@ -13,11 +13,10 @@ def deform_conv2d( offset: Tensor, weight: Tensor, bias: Optional[Tensor] = None, - stride: Union[int, Tuple[int, int]] = (1, 1), - padding: Union[int, Tuple[int, int]] = (0, 0), - dilation: Union[int, Tuple[int, int]] = (1, 1), + stride: Tuple[int, int] = (1, 1), + padding: Tuple[int, int] = (0, 0), + dilation: Tuple[int, int] = (1, 1), ) -> Tensor: - # type: (Tensor, Tensor, Tensor, Optional[Tensor], Tuple[int, int], Tuple[int, int], Tuple[int, int]) -> Tensor """ Performs Deformable Convolution, described in Deformable Convolutional Networks @@ -92,10 +91,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = 1, - padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, + kernel_size: int, + stride: int = 1, + padding: int = 0, + dilation: int = 1, groups: int = 1, bias: bool = True, ): diff --git a/torchvision/ops/feature_pyramid_network.py b/torchvision/ops/feature_pyramid_network.py index a55431f4efd..f0f0b458fa3 100644 --- a/torchvision/ops/feature_pyramid_network.py +++ b/torchvision/ops/feature_pyramid_network.py @@ -72,7 +72,6 @@ def __init__( self.extra_blocks = extra_blocks def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: - # type: (Tensor, int) -> Tensor """ This is equivalent to self.inner_blocks[idx](x), but torchscript doesn't support this yet @@ -91,7 +90,6 @@ def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor: return out def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: - # type: (Tensor, int) -> Tensor """ This is equivalent to self.layer_blocks[idx](x), but torchscript doesn't support this yet @@ -110,7 +108,6 @@ def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor: return out def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: - # type: (Dict[str, Tensor]) -> Dict[str, Tensor] """ Computes the FPN for a set of feature maps. @@ -179,7 +176,6 @@ def forward( y: List[Tensor], names: List[str], ) -> Tuple[List[Tensor], List[str]]: - # type: (List[Tensor], List[Tensor], List[str]) -> Tuple[List[Tensor], List[str]] names.append("pool") x.append(F.max_pool2d(x[-1], 1, 2, 0)) return x, names diff --git a/torchvision/ops/misc.py b/torchvision/ops/misc.py index 6931e325567..a0f15653145 100644 --- a/torchvision/ops/misc.py +++ b/torchvision/ops/misc.py @@ -11,7 +11,7 @@ import warnings import torch from torch import Tensor, Size -from torch.jit.annotations import List, Optional, Tuple, Union +from torch.jit.annotations import List, Optional, Tuple class Conv2d(torch.nn.Conv2d): @@ -50,9 +50,9 @@ class FrozenBatchNorm2d(torch.nn.Module): def __init__( self, - num_features: Union[Size, List[int], Tuple[int, ...]], + num_features: Tuple[int, ...], eps: float = 0., - n: Optional[Size, Union[List[int], Tuple[int, ...]]] = None, + n: Optional[Tuple[int, ...]] = None, ): # n=None for backward-compatibility if n is not None: diff --git a/torchvision/ops/new_empty_tensor.py b/torchvision/ops/new_empty_tensor.py index fb3f2f4ee8c..e964e7a7e15 100644 --- a/torchvision/ops/new_empty_tensor.py +++ b/torchvision/ops/new_empty_tensor.py @@ -4,7 +4,6 @@ def _new_empty_tensor(x: Tensor, shape: List[int]) -> Tensor: - # type: (Tensor, List[int]) -> Tensor """ Arguments: input (Tensor): input tensor diff --git a/torchvision/ops/poolers.py b/torchvision/ops/poolers.py index 7d3f2be7994..adaedb6e4c3 100644 --- a/torchvision/ops/poolers.py +++ b/torchvision/ops/poolers.py @@ -6,7 +6,7 @@ from torchvision.ops import roi_align from torchvision.ops.boxes import box_area -from torch.jit.annotations import Optional, List, Dict, Tuple, Union +from torch.jit.annotations import Optional, List, Dict, Tuple import torchvision @@ -16,7 +16,6 @@ # that merges the levels to the right indices @torch.jit.unused def _onnx_merge_levels(levels: Tensor, unmerged_results: List[Tensor]) -> Tensor: - # type: (Tensor, List[Tensor]) -> Tensor first_result = unmerged_results[0] dtype, device = first_result.dtype, first_result.device res = torch.zeros((levels.size(0), first_result.size(1), @@ -40,7 +39,6 @@ def initLevelMapper( canonical_level: int = 4, eps: float = 1e-6, ): - # type: (int, int, int, int, float) -> LevelMapper return LevelMapper(k_min, k_max, canonical_scale, canonical_level, eps) @@ -64,7 +62,6 @@ def __init__( canonical_level: int = 4, eps: float = 1e-6, ): - # type: (int, int, int, int, float) -> None self.k_min = k_min self.k_max = k_max self.s0 = canonical_scale @@ -72,7 +69,6 @@ def __init__( self.eps = eps def __call__(self, boxlists: List[Tensor]) -> Tensor: - # type: (List[Tensor]) -> Tensor """ Arguments: boxlists (list[BoxList]) @@ -123,7 +119,7 @@ class MultiScaleRoIAlign(nn.Module): def __init__( self, featmap_names: List[str], - output_size: Union[List[Tuple[int, int]], List[int]], + output_size: List[int], sampling_ratio: int, ): super(MultiScaleRoIAlign, self).__init__() @@ -136,7 +132,6 @@ def __init__( self.map_levels = None def convert_to_roi_format(self, boxes: List[Tensor]) -> Tensor: - # type: (List[Tensor]) -> Tensor concat_boxes = torch.cat(boxes, dim=0) device, dtype = concat_boxes.device, concat_boxes.dtype ids = torch.cat( @@ -150,7 +145,6 @@ def convert_to_roi_format(self, boxes: List[Tensor]) -> Tensor: return rois def infer_scale(self, feature: Tensor, original_size: List[int]) -> float: - # type: (Tensor, 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], []) @@ -166,7 +160,6 @@ def setup_scales( features: List[Tensor], image_shapes: List[Tuple[int, int]], ) -> None: - # type: (List[Tensor], List[Tuple[int, int]]) -> None assert len(image_shapes) != 0 max_x = 0 max_y = 0 @@ -186,10 +179,9 @@ def setup_scales( def forward( self, x: Dict[str, Tensor], - boxes: List[Tensor], + boxes: List[Tensor], image_shapes: List[Tuple[int, int]], ) -> Tensor: - # type: (Dict[str, Tensor], List[Tensor], List[Tuple[int, int]]) -> Tensor """ Arguments: x (OrderedDict[Tensor]): feature maps for each level. They are assumed to have diff --git a/torchvision/ops/ps_roi_align.py b/torchvision/ops/ps_roi_align.py index 26ba5bde88c..e3006a3ba53 100644 --- a/torchvision/ops/ps_roi_align.py +++ b/torchvision/ops/ps_roi_align.py @@ -2,19 +2,18 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, Tuple, Union +from torch.jit.annotations import List, Tupl from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def ps_roi_align( input: Tensor, - boxes: Union[Tensor, List[Tensor]], - output_size: Union[int, Tuple[int, int]], + boxes: Tensor, + output_size: int, spatial_scale: float = 1.0, sampling_ratio: int = -1, ) -> Tensor: - # type: (Tensor, Tensor, int, float, int) -> Tensor """ Performs Position-Sensitive Region of Interest (RoI) Align operator mentioned in Light-Head R-CNN. @@ -57,7 +56,7 @@ class PSRoIAlign(nn.Module): """ def __init__( self, - output_size: Union[int, Tuple[int, int]], + output_size: int, spatial_scale: float, sampling_ratio: int, ): @@ -66,7 +65,7 @@ def __init__( self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio - def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + def forward(self, input: Tensor, rois: Tensor) -> Tensor: return ps_roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio) diff --git a/torchvision/ops/ps_roi_pool.py b/torchvision/ops/ps_roi_pool.py index e4b4cf105a5..58c8aa2742a 100644 --- a/torchvision/ops/ps_roi_pool.py +++ b/torchvision/ops/ps_roi_pool.py @@ -2,18 +2,17 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, Tuple, Union +from torch.jit.annotations import List, Tuple from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def ps_roi_pool( input: Tensor, - boxes: Union[Tensor, List[Tensor]], - output_size: Union[int, Tuple[int, int]], + boxes: Tensor, + output_size: int, spatial_scale: float = 1.0, ) -> Tensor: - # type: (Tensor, Tensor, int, float) -> Tensor """ Performs Position-Sensitive Region of Interest (RoI) Pool operator described in R-FCN @@ -48,12 +47,12 @@ class PSRoIPool(nn.Module): """ See ps_roi_pool """ - def __init__(self, output_size: Union[int, Tuple[int, int]], spatial_scale: float): + def __init__(self, output_size: int, spatial_scale: float): super(PSRoIPool, self).__init__() self.output_size = output_size self.spatial_scale = spatial_scale - def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + def forward(self, input: Tensor, rois: Tensor) -> Tensor: return ps_roi_pool(input, rois, self.output_size, self.spatial_scale) def __repr__(self) -> str: diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py index 426e3c37e4d..444f0d7addb 100644 --- a/torchvision/ops/roi_align.py +++ b/torchvision/ops/roi_align.py @@ -2,20 +2,19 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2, Union +from torch.jit.annotations import List, BroadcastingList2 from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def roi_align( input: Tensor, - boxes: Union[Tensor, List[Tensor]], + boxes: Tensor, output_size: BroadcastingList2[int], spatial_scale: float = 1.0, sampling_ratio: int = -1, aligned: bool = False, ) -> Tensor: - # type: (Tensor, Tensor, BroadcastingList2[int], float, int, bool) -> Tensor """ Performs Region of Interest (RoI) Align operator described in Mask R-CNN @@ -69,7 +68,7 @@ def __init__( self.sampling_ratio = sampling_ratio self.aligned = aligned - def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + def forward(self, input: Tensor, rois: Tensor) -> Tensor: return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned) def __repr__(self) -> str: diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py index ec79713cee2..5a71e90d7e7 100644 --- a/torchvision/ops/roi_pool.py +++ b/torchvision/ops/roi_pool.py @@ -2,18 +2,17 @@ from torch import nn, Tensor from torch.nn.modules.utils import _pair -from torch.jit.annotations import List, BroadcastingList2, Union +from torch.jit.annotations import List, BroadcastingList2 from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape def roi_pool( input: Tensor, - boxes: Union[Tensor, List[Tensor]], + boxes: Tensor, output_size: BroadcastingList2[int], spatial_scale: float = 1.0, ) -> Tensor: - # type: (Tensor, Tensor, BroadcastingList2[int], float) -> Tensor """ Performs Region of Interest (RoI) Pool operator described in Fast R-CNN @@ -51,7 +50,7 @@ def __init__(self, output_size: BroadcastingList2[int], spatial_scale: float): self.output_size = output_size self.spatial_scale = spatial_scale - def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor: + def forward(self, input: Tensor, rois: Tensor) -> Tensor: return roi_pool(input, rois, self.output_size, self.spatial_scale) def __repr__(self) -> str: From 0a093cb583b0748706fdea24c3f0448697d7b837 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vuka=C5=A1in=20Manojlovi=C4=87?= Date: Sat, 4 Jul 2020 19:58:31 +0200 Subject: [PATCH 3/4] Fix typo in import --- torchvision/ops/ps_roi_align.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/ops/ps_roi_align.py b/torchvision/ops/ps_roi_align.py index e3006a3ba53..49ee0c21fac 100644 --- a/torchvision/ops/ps_roi_align.py +++ b/torchvision/ops/ps_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, Tupl +from torch.jit.annotations import List, Tuple from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape From c408ac69518035228a74b464f5aa1cccf6f9c77f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vuka=C5=A1in=20Manojlovi=C4=87?= Date: Sat, 4 Jul 2020 21:16:12 +0200 Subject: [PATCH 4/4] Fix undefined name in FeaturePyramidNetwork --- torchvision/ops/feature_pyramid_network.py | 48 +++++++++++----------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/torchvision/ops/feature_pyramid_network.py b/torchvision/ops/feature_pyramid_network.py index f0f0b458fa3..979bbfb1c10 100644 --- a/torchvision/ops/feature_pyramid_network.py +++ b/torchvision/ops/feature_pyramid_network.py @@ -7,6 +7,30 @@ from torch.jit.annotations import Tuple, List, Dict, Optional +class ExtraFPNBlock(nn.Module): + """ + Base class for the extra block in the FPN. + + Arguments: + results (List[Tensor]): the result of the FPN + x (List[Tensor]): the original feature maps + names (List[str]): the names for each one of the + original feature maps + + Returns: + results (List[Tensor]): the extended set of results + of the FPN + names (List[str]): the extended set of names for the results + """ + def forward( + self, + results: List[Tensor], + x: List[Tensor], + names: List[str], + ) -> Tuple[List[Tensor], List[str]]: + pass + + class FeaturePyramidNetwork(nn.Module): """ Module that adds a FPN from on top of a set of feature maps. This is based on @@ -142,30 +166,6 @@ def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: return out -class ExtraFPNBlock(nn.Module): - """ - Base class for the extra block in the FPN. - - Arguments: - results (List[Tensor]): the result of the FPN - x (List[Tensor]): the original feature maps - names (List[str]): the names for each one of the - original feature maps - - Returns: - results (List[Tensor]): the extended set of results - of the FPN - names (List[str]): the extended set of names for the results - """ - def forward( - self, - results: List[Tensor], - x: List[Tensor], - names: List[str], - ) -> Tuple[List[Tensor], List[str]]: - pass - - class LastLevelMaxPool(ExtraFPNBlock): """ Applies a max_pool2d on top of the last feature map