[Enhance] Replace mmdet3d ops with mmcv ops (#1240)

* import some ops from mmcv instead of mmdet3d

* use mmcv ops in primitive_head.py

* use mmcv ops in PAConv

* remove ops in mmdet3d & fix some bugs

* remove spconv & fix some bugs

* fix voxelization unittest

* remove spconv in ops/__init__.py

* refine ops/__init__.py

* recover sparse_block in ops/__init__

* fix parta2_bbox_head unittest

* remove remaining ops

* recover ops/__init__.py for bc breaking

* add source of ops from mmcv

* recover the unittest for voxelization

* remove unittest

mmdet3d/core/bbox/structures/base_box3d.py

@@ -4,9 +4,9 @@ from abc import abstractmethod
 
 import numpy as np
 import torch
+from mmcv._ext import iou3d_boxes_overlap_bev_forward as boxes_overlap_bev_gpu
+from mmcv.ops import points_in_boxes_all, points_in_boxes_part
 
-from mmdet3d.ops import points_in_boxes_all, points_in_boxes_part
-from mmdet3d.ops.iou3d import iou3d_cuda
 from .utils import limit_period, xywhr2xyxyr
 
 
@@ -471,9 +471,8 @@ class BaseInstance3DBoxes(object):
         # bev overlap
         overlaps_bev = boxes1_bev.new_zeros(
             (boxes1_bev.shape[0], boxes2_bev.shape[0])).cuda()  # (N, M)
-        iou3d_cuda.boxes_overlap_bev_gpu(boxes1_bev.contiguous().cuda(),
-                                         boxes2_bev.contiguous().cuda(),
-                                         overlaps_bev)
+        boxes_overlap_bev_gpu(boxes1_bev.contiguous().cuda(),
+                              boxes2_bev.contiguous().cuda(), overlaps_bev)
 
         # 3d overlaps
         overlaps_3d = overlaps_bev.to(boxes1.device) * overlaps_h

mmdet3d/core/post_processing/box3d_nms.py

@@ -2,8 +2,8 @@
 import numba
 import numpy as np
 import torch
-
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 
 
 def box3d_multiclass_nms(mlvl_bboxes,

mmdet3d/core/post_processing/merge_augs.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from ..bbox import bbox3d2result, bbox3d_mapping_back, xywhr2xyxyr
 
 

mmdet3d/models/dense_heads/centerpoint_head.py

@@ -3,6 +3,7 @@ import copy
 
 import torch
 from mmcv.cnn import ConvModule, build_conv_layer
+from mmcv.ops import nms_bev as nms_gpu
 from mmcv.runner import BaseModule, force_fp32
 from torch import nn
 
@@ -11,7 +12,6 @@ from mmdet3d.core import (circle_nms, draw_heatmap_gaussian, gaussian_radius,
 from mmdet3d.models import builder
 from mmdet3d.models.builder import HEADS, build_loss
 from mmdet3d.models.utils import clip_sigmoid
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu
 from mmdet.core import build_bbox_coder, multi_apply
 
 

mmdet3d/models/dense_heads/groupfree3d_head.py

@@ -7,13 +7,14 @@ from mmcv import ConfigDict
 from mmcv.cnn import ConvModule, xavier_init
 from mmcv.cnn.bricks.transformer import (build_positional_encoding,
                                          build_transformer_layer)
+from mmcv.ops import PointsSampler as Points_Sampler
+from mmcv.ops import gather_points
 from mmcv.runner import BaseModule, force_fp32
 from torch import nn as nn
 from torch.nn import functional as F
 
 from mmdet3d.core.post_processing import aligned_3d_nms
 from mmdet3d.models.builder import build_loss
-from mmdet3d.ops import Points_Sampler, gather_points
 from mmdet.core import build_bbox_coder, multi_apply
 from mmdet.models import HEADS
 from .base_conv_bbox_head import BaseConvBboxHead

mmdet3d/models/dense_heads/parta2_rpn_head.py

@@ -3,10 +3,11 @@ from __future__ import division
 
 import numpy as np
 import torch
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 from mmcv.runner import force_fp32
 
 from mmdet3d.core import limit_period, xywhr2xyxyr
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.models import HEADS
 from .anchor3d_head import Anchor3DHead
 

mmdet3d/models/dense_heads/point_rpn_head.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 from mmcv.runner import BaseModule, force_fp32
 from torch import nn as nn
 
 from mmdet3d.core.bbox.structures import (DepthInstance3DBoxes,
                                           LiDARInstance3DBoxes)
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.core import build_bbox_coder, multi_apply
 from mmdet.models import HEADS, build_loss
 

mmdet3d/models/dense_heads/vote_head.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import numpy as np
 import torch
+from mmcv.ops import furthest_point_sample
 from mmcv.runner import BaseModule, force_fp32
 from torch.nn import functional as F
 
@@ -8,7 +9,7 @@ from mmdet3d.core.post_processing import aligned_3d_nms
 from mmdet3d.models.builder import build_loss
 from mmdet3d.models.losses import chamfer_distance
 from mmdet3d.models.model_utils import VoteModule
-from mmdet3d.ops import build_sa_module, furthest_point_sample
+from mmdet3d.ops import build_sa_module
 from mmdet.core import build_bbox_coder, multi_apply
 from mmdet.models import HEADS
 from .base_conv_bbox_head import BaseConvBboxHead

mmdet3d/models/detectors/mvx_two_stage.py

@@ -4,13 +4,13 @@ from os import path as osp
 
 import mmcv
 import torch
+from mmcv.ops import Voxelization
 from mmcv.parallel import DataContainer as DC
 from mmcv.runner import force_fp32
 from torch.nn import functional as F
 
 from mmdet3d.core import (Box3DMode, Coord3DMode, bbox3d2result,
                           merge_aug_bboxes_3d, show_result)
-from mmdet3d.ops import Voxelization
 from mmdet.core import multi_apply
 from mmdet.models import DETECTORS
 from .. import builder

mmdet3d/models/detectors/parta2.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.ops import Voxelization
 from torch.nn import functional as F
 
-from mmdet3d.ops import Voxelization
 from mmdet.models import DETECTORS
 from .. import builder
 from .two_stage import TwoStage3DDetector

mmdet3d/models/detectors/voxelnet.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.ops import Voxelization
 from mmcv.runner import force_fp32
 from torch.nn import functional as F
 
 from mmdet3d.core import bbox3d2result, merge_aug_bboxes_3d
-from mmdet3d.ops import Voxelization
 from mmdet.models import DETECTORS
 from .. import builder
 from .single_stage import SingleStage3DDetector

mmdet3d/models/middle_encoders/sparse_encoder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from mmcv.ops import SparseConvTensor, SparseSequential
 from mmcv.runner import auto_fp16
 from torch import nn as nn
 
 from mmdet3d.ops import SparseBasicBlock, make_sparse_convmodule
-from mmdet3d.ops import spconv as spconv
 from ..builder import MIDDLE_ENCODERS
 
 
@@ -109,9 +109,8 @@ class SparseEncoder(nn.Module):
             dict: Backbone features.
         """
         coors = coors.int()
-        input_sp_tensor = spconv.SparseConvTensor(voxel_features, coors,
-                                                  self.sparse_shape,
-                                                  batch_size)
+        input_sp_tensor = SparseConvTensor(voxel_features, coors,
+                                           self.sparse_shape, batch_size)
         x = self.conv_input(input_sp_tensor)
 
         encode_features = []
@@ -150,7 +149,7 @@ class SparseEncoder(nn.Module):
             int: The number of encoder output channels.
         """
         assert block_type in ['conv_module', 'basicblock']
-        self.encoder_layers = spconv.SparseSequential()
+        self.encoder_layers = SparseSequential()
 
         for i, blocks in enumerate(self.encoder_channels):
             blocks_list = []
@@ -201,6 +200,6 @@ class SparseEncoder(nn.Module):
                             conv_type='SubMConv3d'))
                 in_channels = out_channels
             stage_name = f'encoder_layer{i + 1}'
-            stage_layers = spconv.SparseSequential(*blocks_list)
+            stage_layers = SparseSequential(*blocks_list)
             self.encoder_layers.add_module(stage_name, stage_layers)
         return out_channels

mmdet3d/models/middle_encoders/sparse_unet.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.ops import SparseConvTensor, SparseSequential
 from mmcv.runner import BaseModule, auto_fp16
 
 from mmdet3d.ops import SparseBasicBlock, make_sparse_convmodule
-from mmdet3d.ops import spconv as spconv
 from ..builder import MIDDLE_ENCODERS
 
 
@@ -108,9 +108,8 @@ class SparseUNet(BaseModule):
             dict[str, torch.Tensor]: Backbone features.
         """
         coors = coors.int()
-        input_sp_tensor = spconv.SparseConvTensor(voxel_features, coors,
-                                                  self.sparse_shape,
-                                                  batch_size)
+        input_sp_tensor = SparseConvTensor(voxel_features, coors,
+                                           self.sparse_shape, batch_size)
         x = self.conv_input(input_sp_tensor)
 
         encode_features = []
@@ -200,7 +199,7 @@ class SparseUNet(BaseModule):
         Returns:
             int: The number of encoder output channels.
         """
-        self.encoder_layers = spconv.SparseSequential()
+        self.encoder_layers = SparseSequential()
 
         for i, blocks in enumerate(self.encoder_channels):
             blocks_list = []
@@ -231,7 +230,7 @@ class SparseUNet(BaseModule):
                             conv_type='SubMConv3d'))
                 in_channels = out_channels
             stage_name = f'encoder_layer{i + 1}'
-            stage_layers = spconv.SparseSequential(*blocks_list)
+            stage_layers = SparseSequential(*blocks_list)
             self.encoder_layers.add_module(stage_name, stage_layers)
         return out_channels
 

mmdet3d/models/roi_heads/bbox_heads/parta2_bbox_head.py

@@ -2,6 +2,9 @@
 import numpy as np
 import torch
 from mmcv.cnn import ConvModule, normal_init
+from mmcv.ops import SparseConvTensor, SparseMaxPool3d, SparseSequential
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 from mmcv.runner import BaseModule
 from torch import nn as nn
 
@@ -9,8 +12,6 @@ from mmdet3d.core.bbox.structures import (LiDARInstance3DBoxes,
                                           rotation_3d_in_axis, xywhr2xyxyr)
 from mmdet3d.models.builder import build_loss
 from mmdet3d.ops import make_sparse_convmodule
-from mmdet3d.ops import spconv as spconv
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.core import build_bbox_coder, multi_apply
 from mmdet.models import HEADS
 
@@ -95,7 +96,7 @@ class PartA2BboxHead(BaseModule):
                     indice_key=f'rcnn_part{i}',
                     conv_type='SubMConv3d'))
             part_channel_last = channel
-        self.part_conv = spconv.SparseSequential(*part_conv)
+        self.part_conv = SparseSequential(*part_conv)
 
         seg_channel_last = seg_in_channels
         seg_conv = []
@@ -110,9 +111,9 @@ class PartA2BboxHead(BaseModule):
                     indice_key=f'rcnn_seg{i}',
                     conv_type='SubMConv3d'))
             seg_channel_last = channel
-        self.seg_conv = spconv.SparseSequential(*seg_conv)
+        self.seg_conv = SparseSequential(*seg_conv)
 
-        self.conv_down = spconv.SparseSequential()
+        self.conv_down = SparseSequential()
 
         merge_conv_channel_last = part_channel_last + seg_channel_last
         merge_conv = []
@@ -140,12 +141,10 @@ class PartA2BboxHead(BaseModule):
                     indice_key='rcnn_down1'))
             down_conv_channel_last = channel
 
-        self.conv_down.add_module('merge_conv',
-                                  spconv.SparseSequential(*merge_conv))
-        self.conv_down.add_module(
-            'max_pool3d', spconv.SparseMaxPool3d(kernel_size=2, stride=2))
-        self.conv_down.add_module('down_conv',
-                                  spconv.SparseSequential(*conv_down))
+        self.conv_down.add_module('merge_conv', SparseSequential(*merge_conv))
+        self.conv_down.add_module('max_pool3d',
+                                  SparseMaxPool3d(kernel_size=2, stride=2))
+        self.conv_down.add_module('down_conv', SparseSequential(*conv_down))
 
         shared_fc_list = []
         pool_size = roi_feat_size // 2
@@ -256,10 +255,10 @@ class PartA2BboxHead(BaseModule):
         seg_features = seg_feats[sparse_idx[:, 0], sparse_idx[:, 1],
                                  sparse_idx[:, 2], sparse_idx[:, 3]]
         coords = sparse_idx.int()
-        part_features = spconv.SparseConvTensor(part_features, coords,
-                                                sparse_shape, rcnn_batch_size)
-        seg_features = spconv.SparseConvTensor(seg_features, coords,
-                                               sparse_shape, rcnn_batch_size)
+        part_features = SparseConvTensor(part_features, coords, sparse_shape,
+                                         rcnn_batch_size)
+        seg_features = SparseConvTensor(seg_features, coords, sparse_shape,
+                                        rcnn_batch_size)
 
         # forward rcnn network
         x_part = self.part_conv(part_features)
@@ -267,8 +266,8 @@ class PartA2BboxHead(BaseModule):
 
         merged_feature = torch.cat((x_rpn.features, x_part.features),
                                    dim=1)  # (N, C)
-        shared_feature = spconv.SparseConvTensor(merged_feature, coords,
-                                                 sparse_shape, rcnn_batch_size)
+        shared_feature = SparseConvTensor(merged_feature, coords, sparse_shape,
+                                          rcnn_batch_size)
 
         x = self.conv_down(shared_feature)
 

mmdet3d/models/roi_heads/bbox_heads/point_rcnn_bbox_head.py

@@ -3,6 +3,8 @@ import numpy as np
 import torch
 from mmcv.cnn import ConvModule, normal_init
 from mmcv.cnn.bricks import build_conv_layer
+from mmcv.ops import nms_bev as nms_gpu
+from mmcv.ops import nms_normal_bev as nms_normal_gpu
 from mmcv.runner import BaseModule
 from torch import nn as nn
 
@@ -10,7 +12,6 @@ from mmdet3d.core.bbox.structures import (LiDARInstance3DBoxes,
                                           rotation_3d_in_axis, xywhr2xyxyr)
 from mmdet3d.models.builder import build_loss
 from mmdet3d.ops import build_sa_module
-from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
 from mmdet.core import build_bbox_coder, multi_apply
 from mmdet.models import HEADS
 

mmdet3d/models/roi_heads/mask_heads/primitive_head.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 from mmcv.cnn import ConvModule
+from mmcv.ops import furthest_point_sample
 from mmcv.runner import BaseModule
 from torch import nn as nn
 from torch.nn import functional as F
 
 from mmdet3d.models.builder import build_loss
 from mmdet3d.models.model_utils import VoteModule
-from mmdet3d.ops import build_sa_module, furthest_point_sample
+from mmdet3d.ops import build_sa_module
 from mmdet.core import multi_apply
 from mmdet.models import HEADS
 

mmdet3d/models/roi_heads/roi_extractors/single_roiaware_extractor.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv import ops
 from mmcv.runner import BaseModule
 
-from mmdet3d import ops
 from mmdet.models.builder import ROI_EXTRACTORS
 
 

mmdet3d/models/roi_heads/roi_extractors/single_roipoint_extractor.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv import ops
 from torch import nn as nn
 
-from mmdet3d import ops
 from mmdet3d.core.bbox.structures import rotation_3d_in_axis
 from mmdet.models.builder import ROI_EXTRACTORS
 

mmdet3d/models/voxel_encoders/pillar_encoder.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 from mmcv.cnn import build_norm_layer
+from mmcv.ops import DynamicScatter
 from mmcv.runner import force_fp32
 from torch import nn
 
-from mmdet3d.ops import DynamicScatter
 from ..builder import VOXEL_ENCODERS
 from .utils import PFNLayer, get_paddings_indicator
 

mmdet3d/models/voxel_encoders/voxel_encoder.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 from mmcv.cnn import build_norm_layer
+from mmcv.ops import DynamicScatter
 from mmcv.runner import force_fp32
 from torch import nn
 
-from mmdet3d.ops import DynamicScatter
 from .. import builder
 from ..builder import VOXEL_ENCODERS
 from .utils import VFELayer, get_paddings_indicator