[Feature] Add typehint to mmdet3d/models/backbones.

mmdet3d/models/backbones/base_pointnet.py

@@ -1,14 +1,20 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
 from abc import ABCMeta
+from typing import Optional, Tuple
 
+import torch
 from mmengine.model import BaseModule
 
+from mmdet3d.utils import OptMultiConfig
+
 
 class BasePointNet(BaseModule, metaclass=ABCMeta):
  """Base class for PointNet."""
 
- def __init__(self, init_cfg=None, pretrained=None):
+ def __init__(self,
+ init_cfg: OptMultiConfig = None,
+ pretrained: Optional[str] = None):
  super(BasePointNet, self).__init__(init_cfg)
  self.fp16_enabled = False
  assert not (init_cfg and pretrained), \
@@ -19,7 +25,9 @@ def __init__(self, init_cfg=None, pretrained=None):
  self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
 
  @staticmethod
- def _split_point_feats(points):
+ def _split_point_feats(
+ points: torch.Tensor
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
  """Split coordinates and features of input points.
 
  Args:

mmdet3d/models/backbones/dgcnn.py

@@ -1,9 +1,13 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Sequence, Union
+
+import torch
 from mmengine.model import BaseModule
 from torch import nn as nn
 
 from mmdet3d.models.layers import DGCNNFAModule, DGCNNGFModule
 from mmdet3d.registry import MODELS
+from mmdet3d.utils import ConfigType, OptMultiConfig
 
 
 @MODELS.register_module()
@@ -30,14 +34,15 @@ class DGCNNBackbone(BaseModule):
  """
 
  def __init__(self,
- in_channels,
- num_samples=(20, 20, 20),
- knn_modes=('D-KNN', 'F-KNN', 'F-KNN'),
- radius=(None, None, None),
- gf_channels=((64, 64), (64, 64), (64, )),
- fa_channels=(1024, ),
- act_cfg=dict(type='ReLU'),
- init_cfg=None):
+ in_channels: int,
+ num_samples: Sequence[int] = (20, 20, 20),
+ knn_modes: Sequence[str] = ('D-KNN', 'F-KNN', 'F-KNN'),
+ radius: Sequence[Union[float, None]] = (None, None, None),
+ gf_channels: Sequence[Sequence[int]] = ((64, 64), (64, 64),
+ (64, )),
+ fa_channels: Sequence[int] = (1024, ),
+ act_cfg: ConfigType = dict(type='ReLU'),
+ init_cfg: OptMultiConfig = None):
  super().__init__(init_cfg=init_cfg)
  self.num_gf = len(gf_channels)
 
@@ -71,7 +76,7 @@ def __init__(self,
  self.FA_module = DGCNNFAModule(
  mlp_channels=cur_fa_mlps, act_cfg=act_cfg)
 
- def forward(self, points):
+ def forward(self, points: torch.Tensor) -> dict:
  """Forward pass.
 
  Args:

mmdet3d/models/backbones/dla.py

@@ -1,15 +1,18 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 from mmcv.cnn import build_conv_layer, build_norm_layer
 from mmengine.model import BaseModule
-from torch import nn
+from torch import Tensor, nn
 
 from mmdet3d.registry import MODELS
+from mmdet3d.utils import ConfigType, OptMultiConfig
 
 
-def dla_build_norm_layer(cfg, num_features):
+def dla_build_norm_layer(cfg: ConfigType,
+ num_features: int) -> Tuple[str, nn.Module]:
  """Build normalization layer specially designed for DLANet.
 
  Args:
@@ -53,13 +56,13 @@ class BasicBlock(BaseModule):
  """
 
  def __init__(self,
- in_channels,
- out_channels,
- norm_cfg,
- conv_cfg,
- stride=1,
- dilation=1,
- init_cfg=None):
+ in_channels: int,
+ out_channels: int,
+ norm_cfg: ConfigType,
+ conv_cfg: Dict[str, Any],
+ stride: int = 1,
+ dilation: int = 1,
+ init_cfg: OptMultiConfig = None):
  super(BasicBlock, self).__init__(init_cfg)
  self.conv1 = build_conv_layer(
  conv_cfg,
@@ -84,7 +87,7 @@ def __init__(self,
  self.norm2 = dla_build_norm_layer(norm_cfg, out_channels)[1]
  self.stride = stride
 
- def forward(self, x, identity=None):
+ def forward(self, x: Tensor, identity: Optional[Tensor] = None) -> Tensor:
  """Forward function."""
 
  if identity is None:
@@ -117,13 +120,13 @@ class Root(BaseModule):
  """
 
  def __init__(self,
- in_channels,
- out_channels,
- norm_cfg,
- conv_cfg,
- kernel_size,
- add_identity,
- init_cfg=None):
+ in_channels: int,
+ out_channels: int,
+ norm_cfg: Dict[str, Any],
+ conv_cfg: Dict[str, Any],
+ kernel_size: int,
+ add_identity: bool,
+ init_cfg: OptMultiConfig = None):
  super(Root, self).__init__(init_cfg)
  self.conv = build_conv_layer(
  conv_cfg,
@@ -137,7 +140,7 @@ def __init__(self,
  self.relu = nn.ReLU(inplace=True)
  self.add_identity = add_identity
 
- def forward(self, feat_list):
+ def forward(self, feat_list: List[Tensor]) -> Tensor:
  """Forward function.
 
  Args:
@@ -181,19 +184,19 @@ class Tree(BaseModule):
  """
 
  def __init__(self,
- levels,
- block,
- in_channels,
- out_channels,
- norm_cfg,
- conv_cfg,
- stride=1,
- level_root=False,
- root_dim=None,
- root_kernel_size=1,
- dilation=1,
- add_identity=False,
- init_cfg=None):
+ levels: int,
+ block: nn.Module,
+ in_channels: int,
+ out_channels: int,
+ norm_cfg: dict,
+ conv_cfg: dict,
+ stride: int = 1,
+ level_root: bool = False,
+ root_dim: Optional[int] = None,
+ root_kernel_size: int = 1,
+ dilation: int = 1,
+ add_identity: bool = False,
+ init_cfg: OptMultiConfig = None):
  super(Tree, self).__init__(init_cfg)
  if root_dim is None:
  root_dim = 2 * out_channels
@@ -258,7 +261,10 @@ def __init__(self,
  bias=False),
  dla_build_norm_layer(norm_cfg, out_channels)[1])
 
- def forward(self, x, identity=None, children=None):
+ def forward(self,
+ x: Tensor,
+ identity: Optional[Tensor] = None,
+ children: Optional[List[Tensor]] = None) -> Tensor:
  children = [] if children is None else children
  bottom = self.downsample(x) if self.downsample else x
  identity = self.project(bottom) if self.project else bottom
@@ -302,16 +308,16 @@ class DLANet(BaseModule):
  }
 
  def __init__(self,
- depth,
- in_channels=3,
- out_indices=(0, 1, 2, 3, 4, 5),
- frozen_stages=-1,
- norm_cfg=None,
- conv_cfg=None,
- layer_with_level_root=(False, True, True, True),
- with_identity_root=False,
- pretrained=None,
- init_cfg=None):
+ depth: int,
+ in_channels: int = 3,
+ out_indices: Tuple[int] = (0, 1, 2, 3, 4, 5),
+ frozen_stages: int = -1,
+ norm_cfg: OptMultiConfig = None,
+ conv_cfg: OptMultiConfig = None,
+ layer_with_level_root: List[bool] = (False, True, True, True),
+ with_identity_root: bool = False,
+ pretrained: Optional[str] = None,
+ init_cfg: OptMultiConfig = None):
  super(DLANet, self).__init__(init_cfg)
  if depth not in self.arch_settings:
  raise KeyError(f'invalida depth {depth} for DLA')
@@ -380,13 +386,13 @@ def __init__(self,
  self._freeze_stages()
 
  def _make_conv_level(self,
- in_channels,
- out_channels,
- num_convs,
- norm_cfg,
- conv_cfg,
- stride=1,
- dilation=1):
+ in_channels: int,
+ out_channels: int,
+ num_convs: int,
+ norm_cfg: dict,
+ conv_cfg: dict,
+ stride: int = 1,
+ dilation: int = 1) -> nn.Sequential:
  """Conv modules.
 
  Args:
@@ -418,7 +424,7 @@ def _make_conv_level(self,
  in_channels = out_channels
  return nn.Sequential(*modules)
 
- def _freeze_stages(self):
+ def _freeze_stages(self) -> None:
  if self.frozen_stages >= 0:
  self.base_layer.eval()
  for param in self.base_layer.parameters():
@@ -436,7 +442,7 @@ def _freeze_stages(self):
  for param in m.parameters():
  param.requires_grad = False
 
- def forward(self, x):
+ def forward(self, x: Tensor) -> Tuple[Tensor, ...]:
  outs = []
  x = self.base_layer(x)
  for i in range(self.num_levels):

mmdet3d/models/backbones/multi_backbone.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import copy
 import warnings
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from mmcv.cnn import ConvModule
@@ -27,16 +28,16 @@ class MultiBackbone(BaseModule):
  """
 
  def __init__(self,
- num_streams,
- backbones,
- aggregation_mlp_channels=None,
- conv_cfg=dict(type='Conv1d'),
- norm_cfg=dict(type='BN1d', eps=1e-5, momentum=0.01),
- act_cfg=dict(type='ReLU'),
- suffixes=('net0', 'net1'),
- init_cfg=None,
- pretrained=None,
- **kwargs):
+ num_streams: int,
+ backbones: Union[List, Dict],
+ aggregation_mlp_channels: Optional[List[int]] = None,
+ conv_cfg: Dict = dict(type='Conv1d'),
+ norm_cfg: Dict = dict(type='BN1d', eps=1e-5, momentum=0.01),
+ act_cfg: Dict = dict(type='ReLU'),
+ suffixes: Tuple = ('net0', 'net1'),
+ init_cfg: Optional[Dict] = None,
+ pretrained: Optional[str] = None,
+ **kwargs: Any) -> None:
  super().__init__(init_cfg=init_cfg)
  assert isinstance(backbones, dict) or isinstance(backbones, list)
  if isinstance(backbones, dict):
@@ -89,7 +90,7 @@ def __init__(self,
  'please use "init_cfg" instead')
  self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
 
- def forward(self, points):
+ def forward(self, points: torch.Tensor) -> Dict[str, List[torch.Tensor]]:
  """Forward pass.
 
  Args:

mmdet3d/models/backbones/nostem_regnet.py

@@ -1,4 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
 from mmdet.models.backbones import RegNet
 
 from mmdet3d.registry import MODELS
@@ -59,15 +63,19 @@ class NoStemRegNet(RegNet):
  (1, 1008, 1, 1)
  """
 
- def __init__(self, arch, init_cfg=None, **kwargs):
+ def __init__(self,
+ arch: Dict[str, Union[int, float]],
+ init_cfg: Optional[Dict] = None,
+ **kwargs: Any) -> None:
  super(NoStemRegNet, self).__init__(arch, init_cfg=init_cfg, **kwargs)
 
- def _make_stem_layer(self, in_channels, base_channels):
+ def _make_stem_layer(self, in_channels: int,
+ base_channels: int) -> nn.Module:
  """Override the original function that do not initialize a stem layer
  since 3D detector's voxel encoder works like a stem layer."""
  return
 
- def forward(self, x):
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
  """Forward function of backbone.
 
  Args:

mmdet3d/models/backbones/pointnet2_sa_ssg.py

@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Any, Dict, List, Optional, Tuple
+
 import torch
 from torch import nn as nn
 
@@ -31,20 +33,22 @@ class PointNet2SASSG(BasePointNet):
  """
 
  def __init__(self,
- in_channels,
- num_points=(2048, 1024, 512, 256),
- radius=(0.2, 0.4, 0.8, 1.2),
- num_samples=(64, 32, 16, 16),
- sa_channels=((64, 64, 128), (128, 128, 256), (128, 128, 256),
- (128, 128, 256)),
- fp_channels=((256, 256), (256, 256)),
- norm_cfg=dict(type='BN2d'),
- sa_cfg=dict(
+ in_channels: int,
+ num_points: Tuple[int] = (2048, 1024, 512, 256),
+ radius: Tuple[float] = (0.2, 0.4, 0.8, 1.2),
+ num_samples: Tuple[int] = (64, 32, 16, 16),
+ sa_channels: Tuple[Tuple[int]] = ((64, 64, 128), (128, 128,
+ 256),
+ (128, 128, 256), (128, 128,
+ 256)),
+ fp_channels: Tuple[Tuple[int]] = ((256, 256), (256, 256)),
+ norm_cfg: dict = dict(type='BN2d'),
+ sa_cfg: dict = dict(
  type='PointSAModule',
  pool_mod='max',
  use_xyz=True,
  normalize_xyz=True),
- init_cfg=None):
+ init_cfg: Optional[Dict[str, Any]] = None):
  super().__init__(init_cfg=init_cfg)
  self.num_sa = len(sa_channels)
  self.num_fp = len(fp_channels)
@@ -85,7 +89,7 @@ def __init__(self,
  fp_source_channel = cur_fp_mlps[-1]
  fp_target_channel = skip_channel_list.pop()
 
- def forward(self, points):
+ def forward(self, points: torch.Tensor) -> Dict[str, List[torch.Tensor]]:
  """Forward pass.
 
  Args: