Add type hints for mmcv/ops

mmcv/ops/psa_mask.py

@@ -1,4 +1,7 @@
 # Modified from https://github.com/hszhao/semseg/blob/master/lib/psa
+from typing import Optional, Tuple
+
+import torch
 from torch import nn
 from torch.autograd import Function
 from torch.nn.modules.utils import _pair
@@ -20,7 +23,8 @@ def symbolic(g, input, psa_type, mask_size):
             mask_size_i=mask_size)
 
     @staticmethod
-    def forward(ctx, input, psa_type, mask_size):
+    def forward(ctx, input: torch.Tensor, psa_type: str,
+                mask_size: int) -> torch.Tensor:
         ctx.psa_type = psa_type
         ctx.mask_size = _pair(mask_size)
         ctx.save_for_backward(input)
@@ -45,7 +49,9 @@ def forward(ctx, input, psa_type, mask_size):
         return output
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(
+            ctx, grad_output: torch.Tensor
+    ) -> Tuple[torch.Tensor, None, None, None]:
         input = ctx.saved_tensors[0]
         psa_type = ctx.psa_type
         h_mask, w_mask = ctx.mask_size
@@ -71,7 +77,7 @@ def backward(ctx, grad_output):
 
 class PSAMask(nn.Module):
 
-    def __init__(self, psa_type, mask_size=None):
+    def __init__(self, psa_type: str, mask_size: Optional[tuple] = None):
         super().__init__()
         assert psa_type in ['collect', 'distribute']
         if psa_type == 'collect':
@@ -82,7 +88,7 @@ def __init__(self, psa_type, mask_size=None):
         self.mask_size = mask_size
         self.psa_type = psa_type
 
-    def forward(self, input):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
         return psa_mask(input, self.psa_type_enum, self.mask_size)
 
     def __repr__(self):

mmcv/ops/riroi_align_rotated.py

@@ -1,4 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Optional, Tuple, Union
+
+import torch
 import torch.nn as nn
 from torch.autograd import Function
 
@@ -11,14 +14,14 @@
 class RiRoIAlignRotatedFunction(Function):
 
     @staticmethod
-    def forward(ctx,
-                features,
-                rois,
-                out_size,
-                spatial_scale,
-                num_samples=0,
-                num_orientations=8,
-                clockwise=False):
+    def forward(ctx: Any,
+                features: torch.Tensor,
+                rois: torch.Tensor,
+                out_size: Union[int, tuple],
+                spatial_scale: float,
+                num_samples: int = 0,
+                num_orientations: int = 8,
+                clockwise: bool = False) -> torch.Tensor:
         if isinstance(out_size, int):
             out_h = out_size
             out_w = out_size
@@ -54,7 +57,9 @@ def forward(ctx,
         return output
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(
+        ctx: Any, grad_output: torch.Tensor
+    ) -> Optional[Tuple[torch.Tensor, None, None, None, None, None, None]]:
         feature_size = ctx.feature_size
         spatial_scale = ctx.spatial_scale
         num_orientations = ctx.num_orientations
@@ -67,7 +72,7 @@ def backward(ctx, grad_output):
         out_w = grad_output.size(3)
         out_h = grad_output.size(2)
 
-        grad_input = grad_rois = None
+        grad_input = None
 
         if ctx.needs_input_grad[0]:
             grad_input = rois.new_zeros(batch_size, num_channels, feature_h,
@@ -83,7 +88,8 @@ def backward(ctx, grad_output):
                 num_orientations=num_orientations,
                 clockwise=clockwise)
 
-            return grad_input, grad_rois, None, None, None, None, None
+            return grad_input, None, None, None, None, None, None
+        return None
 
 
 riroi_align_rotated = RiRoIAlignRotatedFunction.apply
@@ -111,11 +117,11 @@ class RiRoIAlignRotated(nn.Module):
     """
 
     def __init__(self,
-                 out_size,
-                 spatial_scale,
-                 num_samples=0,
-                 num_orientations=8,
-                 clockwise=False):
+                 out_size: tuple,
+                 spatial_scale: float,
+                 num_samples: int = 0,
+                 num_orientations: int = 8,
+                 clockwise: bool = False):
         super().__init__()
 
         self.out_size = out_size
@@ -124,7 +130,8 @@ def __init__(self,
         self.num_orientations = int(num_orientations)
         self.clockwise = clockwise
 
-    def forward(self, features, rois):
+    def forward(self, features: torch.Tensor,
+                rois: torch.Tensor) -> torch.Tensor:
         return RiRoIAlignRotatedFunction.apply(features, rois, self.out_size,
                                                self.spatial_scale,
                                                self.num_samples,

mmcv/ops/roi_align.py

@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any
+
 import torch
 import torch.nn as nn
 from torch.autograd import Function
@@ -62,14 +64,14 @@ def symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio,
                 mode_s=pool_mode)
 
     @staticmethod
-    def forward(ctx,
-                input,
-                rois,
-                output_size,
-                spatial_scale=1.0,
-                sampling_ratio=0,
-                pool_mode='avg',
-                aligned=True):
+    def forward(ctx: Any,
+                input: torch.Tensor,
+                rois: torch.Tensor,
+                output_size: int,
+                spatial_scale: float = 1.0,
+                sampling_ratio: int = 0,
+                pool_mode: str = 'avg',
+                aligned: bool = True) -> torch.Tensor:
         ctx.output_size = _pair(output_size)
         ctx.spatial_scale = spatial_scale
         ctx.sampling_ratio = sampling_ratio
@@ -108,7 +110,7 @@ def forward(ctx,
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, grad_output):
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
         rois, argmax_y, argmax_x = ctx.saved_tensors
         grad_input = grad_output.new_zeros(ctx.input_shape)
         # complex head architecture may cause grad_output uncontiguous.
@@ -175,12 +177,12 @@ class RoIAlign(nn.Module):
         },
         cls_name='RoIAlign')
     def __init__(self,
-                 output_size,
-                 spatial_scale=1.0,
-                 sampling_ratio=0,
-                 pool_mode='avg',
-                 aligned=True,
-                 use_torchvision=False):
+                 output_size: tuple,
+                 spatial_scale: float = 1.0,
+                 sampling_ratio: int = 0,
+                 pool_mode: str = 'avg',
+                 aligned: bool = True,
+                 use_torchvision: bool = False):
         super().__init__()
 
         self.output_size = _pair(output_size)
@@ -190,7 +192,7 @@ def __init__(self,
         self.aligned = aligned
         self.use_torchvision = use_torchvision
 
-    def forward(self, input, rois):
+    def forward(self, input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
         """
         Args:
             input: NCHW images

mmcv/ops/roi_align_rotated.py

@@ -1,4 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Optional, Tuple, Union
+
+import torch
 import torch.nn as nn
 from torch.autograd import Function
 from torch.nn.modules.utils import _pair
@@ -37,14 +40,14 @@ def symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio,
             clockwise_i=clockwise)
 
     @staticmethod
-    def forward(ctx,
-                input,
-                rois,
-                output_size,
-                spatial_scale,
-                sampling_ratio=0,
-                aligned=True,
-                clockwise=False):
+    def forward(ctx: Any,
+                input: torch.Tensor,
+                rois: torch.Tensor,
+                output_size: Union[int, tuple],
+                spatial_scale: float,
+                sampling_ratio: int = 0,
+                aligned: bool = True,
+                clockwise: bool = False) -> torch.Tensor:
         ctx.output_size = _pair(output_size)
         ctx.spatial_scale = spatial_scale
         ctx.sampling_ratio = sampling_ratio
@@ -71,7 +74,10 @@ def forward(ctx,
         return output
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(
+        ctx: Any, grad_output: torch.Tensor
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], None, None,
+               None, None, None]:
         feature_size = ctx.feature_size
         rois = ctx.saved_tensors[0]
         assert feature_size is not None
@@ -151,11 +157,11 @@ class RoIAlignRotated(nn.Module):
         },
         cls_name='RoIAlignRotated')
     def __init__(self,
-                 output_size,
-                 spatial_scale,
-                 sampling_ratio=0,
-                 aligned=True,
-                 clockwise=False):
+                 output_size: Union[int, tuple],
+                 spatial_scale: float,
+                 sampling_ratio: int = 0,
+                 aligned: bool = True,
+                 clockwise: bool = False):
         super().__init__()
 
         self.output_size = _pair(output_size)
@@ -164,7 +170,7 @@ def __init__(self,
         self.aligned = aligned
         self.clockwise = clockwise
 
-    def forward(self, input, rois):
+    def forward(self, input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
         return RoIAlignRotatedFunction.apply(input, rois, self.output_size,
                                              self.spatial_scale,
                                              self.sampling_ratio, self.aligned,

mmcv/ops/roi_pool.py

@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Tuple, Union
+
 import torch
 import torch.nn as nn
 from torch.autograd import Function
@@ -23,7 +25,11 @@ def symbolic(g, input, rois, output_size, spatial_scale):
             spatial_scale_f=spatial_scale)
 
     @staticmethod
-    def forward(ctx, input, rois, output_size, spatial_scale=1.0):
+    def forward(ctx: Any,
+                input: torch.Tensor,
+                rois: torch.Tensor,
+                output_size: Union[int, tuple],
+                spatial_scale: float = 1.0) -> torch.Tensor:
         ctx.output_size = _pair(output_size)
         ctx.spatial_scale = spatial_scale
         ctx.input_shape = input.size()
@@ -49,7 +55,9 @@ def forward(ctx, input, rois, output_size, spatial_scale=1.0):
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, grad_output):
+    def backward(
+            ctx: Any, grad_output: torch.Tensor
+    ) -> Tuple[torch.Tensor, None, None, None]:
         rois, argmax = ctx.saved_tensors
         grad_input = grad_output.new_zeros(ctx.input_shape)
 
@@ -70,13 +78,15 @@ def backward(ctx, grad_output):
 
 class RoIPool(nn.Module):
 
-    def __init__(self, output_size, spatial_scale=1.0):
+    def __init__(self,
+                 output_size: Union[int, tuple],
+                 spatial_scale: float = 1.0):
         super().__init__()
 
         self.output_size = _pair(output_size)
         self.spatial_scale = float(spatial_scale)
 
-    def forward(self, input, rois):
+    def forward(self, input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
         return roi_pool(input, rois, self.output_size, self.spatial_scale)
 
     def __repr__(self):

mmcv/ops/roiaware_pool3d.py

@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Tuple, Union
+
 import torch
 from torch import nn as nn
 from torch.autograd import Function
@@ -25,7 +27,10 @@ class RoIAwarePool3d(nn.Module):
             Default: 'max'.
     """
 
-    def __init__(self, out_size, max_pts_per_voxel=128, mode='max'):
+    def __init__(self,
+                 out_size: Union[int, tuple],
+                 max_pts_per_voxel: int = 128,
+                 mode: str = 'max'):
         super().__init__()
 
         self.out_size = out_size
@@ -34,7 +39,8 @@ def __init__(self, out_size, max_pts_per_voxel=128, mode='max'):
         pool_mapping = {'max': 0, 'avg': 1}
         self.mode = pool_mapping[mode]
 
-    def forward(self, rois, pts, pts_feature):
+    def forward(self, rois: torch.Tensor, pts: torch.Tensor,
+                pts_feature: torch.Tensor) -> torch.Tensor:
         """
         Args:
             rois (torch.Tensor): [N, 7], in LiDAR coordinate,
@@ -55,8 +61,9 @@ def forward(self, rois, pts, pts_feature):
 class RoIAwarePool3dFunction(Function):
 
     @staticmethod
-    def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel,
-                mode):
+    def forward(ctx: Any, rois: torch.Tensor, pts: torch.Tensor,
+                pts_feature: torch.Tensor, out_size: Union[int, tuple],
+                max_pts_per_voxel: int, mode: int) -> torch.Tensor:
         """
         Args:
             rois (torch.Tensor): [N, 7], in LiDAR coordinate,
@@ -108,7 +115,9 @@ def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel,
         return pooled_features
 
     @staticmethod
-    def backward(ctx, grad_out):
+    def backward(
+        ctx: Any, grad_out: torch.Tensor
+    ) -> Tuple[None, None, torch.Tensor, None, None, None]:
         ret = ctx.roiaware_pool3d_for_backward
         pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret