[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

tests/test_models/test_common_modules/test_sparse_unet.py

 # Copyright (c) OpenMMLab. All rights reserved.
+import pytest
 import torch
+from mmcv.ops import (SparseConv3d, SparseConvTensor, SparseInverseConv3d,
+                      SubMConv3d)
 
 from mmdet3d.ops import SparseBasicBlock
-from mmdet3d.ops import spconv as spconv
 
 
 def test_SparseUNet():
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and torch+cuda')
     from mmdet3d.models.middle_encoders.sparse_unet import SparseUNet
-    self = SparseUNet(in_channels=4, sparse_shape=[41, 1600, 1408])
+    self = SparseUNet(in_channels=4, sparse_shape=[41, 1600, 1408]).cuda()
 
     # test encoder layers
     assert len(self.encoder_layers) == 4
     assert self.encoder_layers.encoder_layer1[0][0].in_channels == 16
     assert self.encoder_layers.encoder_layer1[0][0].out_channels == 16
-    assert isinstance(self.encoder_layers.encoder_layer1[0][0],
-                      spconv.conv.SubMConv3d)
+    assert isinstance(self.encoder_layers.encoder_layer1[0][0], SubMConv3d)
     assert isinstance(self.encoder_layers.encoder_layer1[0][1],
                       torch.nn.modules.batchnorm.BatchNorm1d)
     assert isinstance(self.encoder_layers.encoder_layer1[0][2],
                       torch.nn.modules.activation.ReLU)
     assert self.encoder_layers.encoder_layer4[0][0].in_channels == 64
     assert self.encoder_layers.encoder_layer4[0][0].out_channels == 64
-    assert isinstance(self.encoder_layers.encoder_layer4[0][0],
-                      spconv.conv.SparseConv3d)
-    assert isinstance(self.encoder_layers.encoder_layer4[2][0],
-                      spconv.conv.SubMConv3d)
+    assert isinstance(self.encoder_layers.encoder_layer4[0][0], SparseConv3d)
+    assert isinstance(self.encoder_layers.encoder_layer4[2][0], SubMConv3d)
 
     # test decoder layers
     assert isinstance(self.lateral_layer1, SparseBasicBlock)
-    assert isinstance(self.merge_layer1[0], spconv.conv.SubMConv3d)
-    assert isinstance(self.upsample_layer1[0], spconv.conv.SubMConv3d)
-    assert isinstance(self.upsample_layer2[0], spconv.conv.SparseInverseConv3d)
-
-    voxel_features = torch.tensor([[6.56126, 0.9648336, -1.7339306, 0.315],
-                                   [6.8162713, -2.480431, -1.3616394, 0.36],
-                                   [11.643568, -4.744306, -1.3580885, 0.16],
-                                   [23.482342, 6.5036807, 0.5806964, 0.35]],
-                                  dtype=torch.float32)  # n, point_features
+    assert isinstance(self.merge_layer1[0], SubMConv3d)
+    assert isinstance(self.upsample_layer1[0], SubMConv3d)
+    assert isinstance(self.upsample_layer2[0], SparseInverseConv3d)
+
+    voxel_features = torch.tensor(
+        [[6.56126, 0.9648336, -1.7339306, 0.315],
+         [6.8162713, -2.480431, -1.3616394, 0.36],
+         [11.643568, -4.744306, -1.3580885, 0.16],
+         [23.482342, 6.5036807, 0.5806964, 0.35]],
+        dtype=torch.float32).cuda()  # n, point_features
     coordinates = torch.tensor(
         [[0, 12, 819, 131], [0, 16, 750, 136], [1, 16, 705, 232],
          [1, 35, 930, 469]],
-        dtype=torch.int32)  # n, 4(batch, ind_x, ind_y, ind_z)
+        dtype=torch.int32).cuda()  # n, 4(batch, ind_x, ind_y, ind_z)
 
     unet_ret_dict = self.forward(voxel_features, coordinates, 2)
     seg_features = unet_ret_dict['seg_features']
@@ -51,29 +53,32 @@ def test_SparseUNet():
 
 
 def test_SparseBasicBlock():
-    voxel_features = torch.tensor([[6.56126, 0.9648336, -1.7339306, 0.315],
-                                   [6.8162713, -2.480431, -1.3616394, 0.36],
-                                   [11.643568, -4.744306, -1.3580885, 0.16],
-                                   [23.482342, 6.5036807, 0.5806964, 0.35]],
-                                  dtype=torch.float32)  # n, point_features
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and torch+cuda')
+    voxel_features = torch.tensor(
+        [[6.56126, 0.9648336, -1.7339306, 0.315],
+         [6.8162713, -2.480431, -1.3616394, 0.36],
+         [11.643568, -4.744306, -1.3580885, 0.16],
+         [23.482342, 6.5036807, 0.5806964, 0.35]],
+        dtype=torch.float32).cuda()  # n, point_features
     coordinates = torch.tensor(
         [[0, 12, 819, 131], [0, 16, 750, 136], [1, 16, 705, 232],
          [1, 35, 930, 469]],
-        dtype=torch.int32)  # n, 4(batch, ind_x, ind_y, ind_z)
+        dtype=torch.int32).cuda()  # n, 4(batch, ind_x, ind_y, ind_z)
 
     # test
-    input_sp_tensor = spconv.SparseConvTensor(voxel_features, coordinates,
-                                              [41, 1600, 1408], 2)
+    input_sp_tensor = SparseConvTensor(voxel_features, coordinates,
+                                       [41, 1600, 1408], 2)
     self = SparseBasicBlock(
         4,
         4,
         conv_cfg=dict(type='SubMConv3d', indice_key='subm1'),
-        norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01))
+        norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01)).cuda()
     # test conv and bn layer
-    assert isinstance(self.conv1, spconv.conv.SubMConv3d)
+    assert isinstance(self.conv1, SubMConv3d)
     assert self.conv1.in_channels == 4
     assert self.conv1.out_channels == 4
-    assert isinstance(self.conv2, spconv.conv.SubMConv3d)
+    assert isinstance(self.conv2, SubMConv3d)
     assert self.conv2.out_channels == 4
     assert self.conv2.out_channels == 4
     assert self.bn1.eps == 1e-3
@@ -84,21 +89,24 @@ def test_SparseBasicBlock():
 
 
 def test_make_sparse_convmodule():
+    if not torch.cuda.is_available():
+        pytest.skip('test requires GPU and torch+cuda')
     from mmdet3d.ops import make_sparse_convmodule
 
-    voxel_features = torch.tensor([[6.56126, 0.9648336, -1.7339306, 0.315],
-                                   [6.8162713, -2.480431, -1.3616394, 0.36],
-                                   [11.643568, -4.744306, -1.3580885, 0.16],
-                                   [23.482342, 6.5036807, 0.5806964, 0.35]],
-                                  dtype=torch.float32)  # n, point_features
+    voxel_features = torch.tensor(
+        [[6.56126, 0.9648336, -1.7339306, 0.315],
+         [6.8162713, -2.480431, -1.3616394, 0.36],
+         [11.643568, -4.744306, -1.3580885, 0.16],
+         [23.482342, 6.5036807, 0.5806964, 0.35]],
+        dtype=torch.float32).cuda()  # n, point_features
     coordinates = torch.tensor(
         [[0, 12, 819, 131], [0, 16, 750, 136], [1, 16, 705, 232],
          [1, 35, 930, 469]],
-        dtype=torch.int32)  # n, 4(batch, ind_x, ind_y, ind_z)
+        dtype=torch.int32).cuda()  # n, 4(batch, ind_x, ind_y, ind_z)
 
     # test
-    input_sp_tensor = spconv.SparseConvTensor(voxel_features, coordinates,
-                                              [41, 1600, 1408], 2)
+    input_sp_tensor = SparseConvTensor(voxel_features, coordinates,
+                                       [41, 1600, 1408], 2)
 
     sparse_block0 = make_sparse_convmodule(
         4,
@@ -109,8 +117,8 @@ def test_make_sparse_convmodule():
         padding=0,
         conv_type='SubMConv3d',
         norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
-        order=('conv', 'norm', 'act'))
-    assert isinstance(sparse_block0[0], spconv.SubMConv3d)
+        order=('conv', 'norm', 'act')).cuda()
+    assert isinstance(sparse_block0[0], SubMConv3d)
     assert sparse_block0[0].in_channels == 4
     assert sparse_block0[0].out_channels == 16
     assert isinstance(sparse_block0[1], torch.nn.BatchNorm1d)
@@ -134,4 +142,4 @@ def test_make_sparse_convmodule():
         order=('norm', 'act', 'conv'))
     assert isinstance(sparse_block1[0], torch.nn.BatchNorm1d)
     assert isinstance(sparse_block1[1], torch.nn.ReLU)
-    assert isinstance(sparse_block1[2], spconv.SparseInverseConv3d)
+    assert isinstance(sparse_block1[2], SparseInverseConv3d)

tests/test_models/test_heads/test_parta2_bbox_head.py

@@ -2,13 +2,13 @@
 import pytest
 import torch
 from mmcv import Config
+from mmcv.ops import SubMConv3d
 from torch.nn import BatchNorm1d, ReLU
 
 from mmdet3d.core.bbox import Box3DMode, LiDARInstance3DBoxes
 from mmdet3d.core.bbox.samplers import IoUNegPiecewiseSampler
 from mmdet3d.models import PartA2BboxHead
 from mmdet3d.ops import make_sparse_convmodule
-from mmdet3d.ops.spconv.conv import SubMConv3d
 
 
 def test_loss():

tests/test_models/test_voxel_encoder/test_dynamic_scatter.py

-# Copyright (c) OpenMMLab. All rights reserved.
-import pytest
-import torch
-from torch.autograd import gradcheck
-
-from mmdet3d.ops import DynamicScatter
-
-
-def test_dynamic_scatter():
-    if not torch.cuda.is_available():
-        pytest.skip('test requires GPU and torch+cuda')
-
-    dsmean = DynamicScatter([0.32, 0.32, 6],
-                            [-74.88, -74.88, -2, 74.88, 74.88, 4], True)
-    dsmax = DynamicScatter([0.32, 0.32, 6],
-                           [-74.88, -74.88, -2, 74.88, 74.88, 4], False)
-
-    # test empty input
-    empty_feats = torch.empty(size=(0, 3), dtype=torch.float32, device='cuda')
-    empty_coors = torch.empty(size=(0, 3), dtype=torch.int32, device='cuda')
-
-    empty_feats.requires_grad_()
-    empty_feats_out_mean, empty_coors_out_mean = dsmean(
-        empty_feats, empty_coors)
-    empty_feats_out_mean.sum().backward()
-    empty_feats_out_max, empty_coors_out_max = dsmax(empty_feats, empty_coors)
-    empty_feats_out_max.sum().backward()
-
-    assert empty_feats_out_mean.shape == empty_feats.shape
-    assert empty_feats_out_max.shape == empty_feats.shape
-    assert empty_coors_out_mean.shape == empty_coors.shape
-    assert empty_coors_out_max.shape == empty_coors.shape
-
-    # test empty reduced output
-    empty_o_feats = torch.rand(
-        size=(200000, 3), dtype=torch.float32, device='cuda') * 100 - 50
-    empty_o_coors = torch.randint(
-        low=-1, high=0, size=(200000, 3), dtype=torch.int32, device='cuda')
-
-    empty_o_feats.requires_grad_()
-    empty_o_feats_out_mean, empty_o_coors_out_mean = dsmean(
-        empty_o_feats, empty_o_coors)
-    empty_o_feats_out_mean.sum().backward()
-    assert (empty_o_feats.grad == 0).all()
-
-    empty_o_feats_out_max, empty_o_coors_out_max = dsmax(
-        empty_o_feats, empty_o_coors)
-    empty_o_feats_out_max.sum().backward()
-    assert (empty_o_feats.grad == 0).all()
-
-    # test non-empty input
-    feats = torch.rand(
-        size=(200000, 3), dtype=torch.float32, device='cuda') * 100 - 50
-    coors = torch.randint(
-        low=-1, high=20, size=(200000, 3), dtype=torch.int32, device='cuda')
-
-    ref_voxel_coors = coors.unique(dim=0, sorted=True)
-    ref_voxel_coors = ref_voxel_coors[ref_voxel_coors.min(dim=-1).values >= 0]
-    ref_voxel_feats_mean = []
-    ref_voxel_feats_max = []
-    for ref_voxel_coor in ref_voxel_coors:
-        voxel_mask = (coors == ref_voxel_coor).all(dim=-1)
-        ref_voxel_feats_mean.append(feats[voxel_mask].mean(dim=0))
-        ref_voxel_feats_max.append(feats[voxel_mask].max(dim=0).values)
-    ref_voxel_feats_mean = torch.stack(ref_voxel_feats_mean)
-    ref_voxel_feats_max = torch.stack(ref_voxel_feats_max)
-
-    feats_out_mean, coors_out_mean = dsmean(feats, coors)
-    seq_mean = (coors_out_mean[:, 0] * 400 + coors_out_mean[:, 1] * 20 +
-                coors_out_mean[:, 2]).argsort()
-    feats_out_mean = feats_out_mean[seq_mean]
-    coors_out_mean = coors_out_mean[seq_mean]
-
-    feats_out_max, coors_out_max = dsmax(feats, coors)
-    seq_max = (coors_out_max[:, 0] * 400 + coors_out_max[:, 1] * 20 +
-               coors_out_max[:, 2]).argsort()
-    feats_out_max = feats_out_max[seq_max]
-    coors_cout_max = coors_out_max[seq_max]
-
-    assert (coors_out_mean == ref_voxel_coors).all()
-    assert torch.allclose(
-        feats_out_mean, ref_voxel_feats_mean, atol=1e-2, rtol=1e-5)
-    assert (coors_cout_max == ref_voxel_coors).all()
-    assert torch.allclose(
-        feats_out_max, ref_voxel_feats_max, atol=1e-2, rtol=1e-5)
-
-    # test non-empty input without any point out of bound
-    feats = torch.rand(
-        size=(200000, 3), dtype=torch.float32, device='cuda') * 100 - 50
-    coors = torch.randint(
-        low=0, high=20, size=(200000, 3), dtype=torch.int32, device='cuda')
-
-    ref_voxel_coors = coors.unique(dim=0, sorted=True)
-    ref_voxel_coors = ref_voxel_coors[ref_voxel_coors.min(dim=-1).values >= 0]
-    ref_voxel_feats_mean = []
-    ref_voxel_feats_max = []
-    for ref_voxel_coor in ref_voxel_coors:
-        voxel_mask = (coors == ref_voxel_coor).all(dim=-1)
-        ref_voxel_feats_mean.append(feats[voxel_mask].mean(dim=0))
-        ref_voxel_feats_max.append(feats[voxel_mask].max(dim=0).values)
-    ref_voxel_feats_mean = torch.stack(ref_voxel_feats_mean)
-    ref_voxel_feats_max = torch.stack(ref_voxel_feats_max)
-
-    feats_out_mean, coors_out_mean = dsmean(feats, coors)
-    seq_mean = (coors_out_mean[:, 0] * 400 + coors_out_mean[:, 1] * 20 +
-                coors_out_mean[:, 2]).argsort()
-    feats_out_mean = feats_out_mean[seq_mean]
-    coors_out_mean = coors_out_mean[seq_mean]
-
-    feats_out_max, coors_out_max = dsmax(feats, coors)
-    seq_max = (coors_out_max[:, 0] * 400 + coors_out_max[:, 1] * 20 +
-               coors_out_max[:, 2]).argsort()
-    feats_out_max = feats_out_max[seq_max]
-    coors_cout_max = coors_out_max[seq_max]
-
-    assert (coors_out_mean == ref_voxel_coors).all()
-    assert torch.allclose(
-        feats_out_mean, ref_voxel_feats_mean, atol=1e-2, rtol=1e-5)
-    assert (coors_cout_max == ref_voxel_coors).all()
-    assert torch.allclose(
-        feats_out_max, ref_voxel_feats_max, atol=1e-2, rtol=1e-5)
-
-    # test grad #
-    feats = torch.rand(
-        size=(100, 4), dtype=torch.float32, device='cuda') * 100 - 50
-    coors = torch.randint(
-        low=-1, high=3, size=(100, 3), dtype=torch.int32, device='cuda')
-    feats.requires_grad_()
-    gradcheck(dsmean, (feats, coors), eps=1e-2, atol=1e-2, rtol=1e-5)
-    gradcheck(dsmax, (feats, coors), eps=1e-2, atol=1e-2, rtol=1e-5)

tests/test_models/test_voxel_encoder/test_voxelize.py

-# Copyright (c) OpenMMLab. All rights reserved.
-import numpy as np
-import pytest
-import torch
-
-from mmdet3d.core.voxel.voxel_generator import VoxelGenerator
-from mmdet3d.datasets.pipelines import LoadPointsFromFile
-from mmdet3d.ops.voxel.voxelize import Voxelization
-
-
-def _get_voxel_points_indices(points, coors, voxel):
-    result_form = np.equal(coors, voxel)
-    return result_form[:, 0] & result_form[:, 1] & result_form[:, 2]
-
-
-def test_voxelization():
-    voxel_size = [0.5, 0.5, 0.5]
-    point_cloud_range = [0, -40, -3, 70.4, 40, 1]
-    max_num_points = 1000
-    self = VoxelGenerator(voxel_size, point_cloud_range, max_num_points)
-    data_path = './tests/data/kitti/training/velodyne_reduced/000000.bin'
-    load_points_from_file = LoadPointsFromFile(
-        coord_type='LIDAR', load_dim=4, use_dim=4)
-    results = dict()
-    results['pts_filename'] = data_path
-    results = load_points_from_file(results)
-    points = results['points'].tensor.numpy()
-    voxels_generator = self.generate(points)
-    coors, voxels, num_points_per_voxel = voxels_generator
-    expected_coors = coors
-    expected_voxels = voxels
-    expected_num_points_per_voxel = num_points_per_voxel
-
-    points = torch.tensor(points)
-    max_num_points = -1
-    dynamic_voxelization = Voxelization(voxel_size, point_cloud_range,
-                                        max_num_points)
-    max_num_points = 1000
-    hard_voxelization = Voxelization(voxel_size, point_cloud_range,
-                                     max_num_points)
-    # test hard_voxelization on cpu
-    coors, voxels, num_points_per_voxel = hard_voxelization.forward(points)
-    coors = coors.detach().numpy()
-    voxels = voxels.detach().numpy()
-    num_points_per_voxel = num_points_per_voxel.detach().numpy()
-    assert np.all(coors == expected_coors)
-    assert np.all(voxels == expected_voxels)
-    assert np.all(num_points_per_voxel == expected_num_points_per_voxel)
-
-    # test dynamic_voxelization on cpu
-    coors = dynamic_voxelization.forward(points)
-    coors = coors.detach().numpy()
-    points = points.detach().numpy()
-    for i in range(expected_voxels.shape[0]):
-        indices = _get_voxel_points_indices(points, coors, expected_voxels[i])
-        num_points_current_voxel = points[indices].shape[0]
-        assert num_points_current_voxel > 0
-        assert np.all(
-            points[indices] == expected_coors[i][:num_points_current_voxel])
-        assert num_points_current_voxel == expected_num_points_per_voxel[i]
-
-    if not torch.cuda.is_available():
-        pytest.skip('test requires GPU and torch+cuda')
-    # test hard_voxelization on gpu
-    points = torch.tensor(points).contiguous().to(device='cuda:0')
-    coors, voxels, num_points_per_voxel = hard_voxelization.forward(points)
-    coors = coors.cpu().detach().numpy()
-    voxels = voxels.cpu().detach().numpy()
-    num_points_per_voxel = num_points_per_voxel.cpu().detach().numpy()
-    assert np.all(coors == expected_coors)
-    assert np.all(voxels == expected_voxels)
-    assert np.all(num_points_per_voxel == expected_num_points_per_voxel)
-
-    # test dynamic_voxelization on gpu
-    coors = dynamic_voxelization.forward(points)
-    coors = coors.cpu().detach().numpy()
-    points = points.cpu().detach().numpy()
-    for i in range(expected_voxels.shape[0]):
-        indices = _get_voxel_points_indices(points, coors, expected_voxels[i])
-        num_points_current_voxel = points[indices].shape[0]
-        assert num_points_current_voxel > 0
-        assert np.all(
-            points[indices] == expected_coors[i][:num_points_current_voxel])
-        assert num_points_current_voxel == expected_num_points_per_voxel[i]
-
-
-def test_voxelization_nondeterministic():
-    if not torch.cuda.is_available():
-        pytest.skip('test requires GPU and torch+cuda')
-
-    voxel_size = [0.5, 0.5, 0.5]
-    point_cloud_range = [0, -40, -3, 70.4, 40, 1]
-    data_path = './tests/data/kitti/training/velodyne_reduced/000000.bin'
-    load_points_from_file = LoadPointsFromFile(
-        coord_type='LIDAR', load_dim=4, use_dim=4)
-    results = dict()
-    results['pts_filename'] = data_path
-    results = load_points_from_file(results)
-    points = results['points'].tensor.numpy()
-
-    points = torch.tensor(points)
-    max_num_points = -1
-    dynamic_voxelization = Voxelization(voxel_size, point_cloud_range,
-                                        max_num_points)
-
-    max_num_points = 10
-    max_voxels = 50
-    hard_voxelization = Voxelization(
-        voxel_size,
-        point_cloud_range,
-        max_num_points,
-        max_voxels,
-        deterministic=False)
-
-    # test hard_voxelization (non-deterministic version) on gpu
-    points = torch.tensor(points).contiguous().to(device='cuda:0')
-    voxels, coors, num_points_per_voxel = hard_voxelization.forward(points)
-    coors = coors.cpu().detach().numpy().tolist()
-    voxels = voxels.cpu().detach().numpy().tolist()
-    num_points_per_voxel = num_points_per_voxel.cpu().detach().numpy().tolist()
-
-    coors_all = dynamic_voxelization.forward(points)
-    coors_all = coors_all.cpu().detach().numpy().tolist()
-
-    coors_set = set([tuple(c) for c in coors])
-    coors_all_set = set([tuple(c) for c in coors_all])
-
-    assert len(coors_set) == len(coors)
-    assert len(coors_set - coors_all_set) == 0
-
-    points = points.cpu().detach().numpy().tolist()
-
-    coors_points_dict = {}
-    for c, ps in zip(coors_all, points):
-        if tuple(c) not in coors_points_dict:
-            coors_points_dict[tuple(c)] = set()
-        coors_points_dict[tuple(c)].add(tuple(ps))
-
-    for c, ps, n in zip(coors, voxels, num_points_per_voxel):
-        ideal_voxel_points_set = coors_points_dict[tuple(c)]
-        voxel_points_set = set([tuple(p) for p in ps[:n]])
-        assert len(voxel_points_set) == n
-        if n < max_num_points:
-            assert voxel_points_set == ideal_voxel_points_set
-            for p in ps[n:]:
-                assert max(p) == min(p) == 0
-        else:
-            assert len(voxel_points_set - ideal_voxel_points_set) == 0
-
-    # test hard_voxelization (non-deterministic version) on gpu
-    # with all input point in range
-    points = torch.tensor(points).contiguous().to(device='cuda:0')[:max_voxels]
-    coors_all = dynamic_voxelization.forward(points)
-    valid_mask = coors_all.ge(0).all(-1)
-    points = points[valid_mask]
-    coors_all = coors_all[valid_mask]
-    coors_all = coors_all.cpu().detach().numpy().tolist()
-
-    voxels, coors, num_points_per_voxel = hard_voxelization.forward(points)
-    coors = coors.cpu().detach().numpy().tolist()
-
-    coors_set = set([tuple(c) for c in coors])
-    coors_all_set = set([tuple(c) for c in coors_all])
-
-    assert len(coors_set) == len(coors) == len(coors_all_set)