[Enhance] GroupFree3d inherits BaseModule From MMCV (#704)

mmdet3d/models/dense_heads/groupfree3d_head.py

@@ -2,10 +2,10 @@
 import numpy as np
 import torch
 from mmcv import ConfigDict
-from mmcv.cnn import ConvModule
+from mmcv.cnn import ConvModule, xavier_init
 from mmcv.cnn.bricks.transformer import (build_positional_encoding,
  build_transformer_layer)
-from mmcv.runner import force_fp32
+from mmcv.runner import BaseModule, force_fp32
 from torch import nn as nn
 from torch.nn import functional as F
 
@@ -19,7 +19,7 @@
 EPS = 1e-6
 
 
-class PointsObjClsModule(nn.Module):
+class PointsObjClsModule(BaseModule):
  """object candidate point prediction from seed point features.
 
  Args:
@@ -39,8 +39,9 @@ def __init__(self,
  num_convs=3,
  conv_cfg=dict(type='Conv1d'),
  norm_cfg=dict(type='BN1d'),
- act_cfg=dict(type='ReLU')):
- super().__init__()
+ act_cfg=dict(type='ReLU'),
+ init_cfg=None):
+ super().__init__(init_cfg=init_cfg)
  conv_channels = [in_channel for _ in range(num_convs - 1)]
  conv_channels.append(1)
 
@@ -104,7 +105,7 @@ def forward(self, xyz, features, sample_inds):
 
 
 @HEADS.register_module()
-class GroupFree3DHead(nn.Module):
+class GroupFree3DHead(BaseModule):
  r"""Bbox head of `Group-Free 3D <https://arxiv.org/abs/2104.00678>`_.
 
  Args:
@@ -162,8 +163,9 @@ def __init__(self,
  size_class_loss=None,
  size_res_loss=None,
  size_reg_loss=None,
- semantic_loss=None):
- super(GroupFree3DHead, self).__init__()
+ semantic_loss=None,
+ init_cfg=None):
+ super(GroupFree3DHead, self).__init__(init_cfg=init_cfg)
  self.num_classes = num_classes
  self.train_cfg = train_cfg
  self.test_cfg = test_cfg
@@ -251,15 +253,13 @@ def init_weights(self):
  # initialize transformer
  for m in self.decoder_layers.parameters():
  if m.dim() > 1:
- nn.init.xavier_uniform_(m)
-
+ xavier_init(m, distribution='uniform')
  for m in self.decoder_self_posembeds.parameters():
  if m.dim() > 1:
- nn.init.xavier_uniform_(m)
-
+ xavier_init(m, distribution='uniform')
  for m in self.decoder_cross_posembeds.parameters():
  if m.dim() > 1:
- nn.init.xavier_uniform_(m)
+ xavier_init(m, distribution='uniform')
 
  def _get_cls_out_channels(self):
  """Return the channel number of classification outputs."""