refactor voxel-based segmentor

configs/_base_/models/cylinder3d.py

@@ -1,7 +1,7 @@
 grid_shape = [480, 360, 32]
 point_cloud_range = [0, -3.14159265359, -4, 50, 3.14159265359, 2]
 model = dict(
-    type='Cylinder3D',
+    type='VoxelSegmentor',
     data_preprocessor=dict(
         type='Det3DDataPreprocessor',
         voxel=True,

configs/_base_/models/minkunet.py

@@ -1,5 +1,5 @@
 model = dict(
-    type='MinkUNet',
+    type='VoxelSegmentor',
     data_preprocessor=dict(
         type='Det3DDataPreprocessor',
         voxel=True,
@@ -26,6 +26,7 @@
         type='MinkUNetHead',
         channels=96,
         num_classes=19,
+        batch_first=False,
         dropout_ratio=0,
         loss_ce=dict(type='mmdet.CrossEntropyLoss', avg_non_ignore=True),
         ignore_index=19),

configs/_base_/models/spvcnn.py

@@ -1,5 +1,5 @@
 model = dict(
-    type='MinkUNet',
+    type='VoxelSegmentor',
     data_preprocessor=dict(
         type='Det3DDataPreprocessor',
         voxel=True,
@@ -27,6 +27,7 @@
         type='MinkUNetHead',
         channels=96,
         num_classes=19,
+        batch_first=False,
         dropout_ratio=0,
         loss_ce=dict(type='mmdet.CrossEntropyLoss', avg_non_ignore=True),
         ignore_index=19),

configs/minkunet/minkunet18_w16_torchsparse_8xb2-amp-15e_semantickitti.py

@@ -1,4 +1,4 @@
-_base_ = ['./minkunet_w32_8xb2-15e_semantickitti.py']
+_base_ = ['./minkunet18_w32_torchsparse_8xb2-amp-15e_semantickitti.py']
 
 model = dict(
     backbone=dict(

configs/minkunet/minkunet18_w20_torchsparse_8xb2-amp-15e_semantickitti.py

@@ -1,4 +1,4 @@
-_base_ = ['./minkunet_w32_8xb2-15e_semantickitti.py']
+_base_ = ['./minkunet18_w32_torchsparse_8xb2-amp-15e_semantickitti.py']
 
 model = dict(
     backbone=dict(

configs/minkunet/minkunet34_w32_minkowski_8xb2-laser-polar-mix-3x_semantickitti.py

@@ -4,4 +4,5 @@
 
 model = dict(
     data_preprocessor=dict(batch_first=True),
-    backbone=dict(sparseconv_backend='minkowski'))
+    backbone=dict(sparseconv_backend='minkowski'),
+    decode_head=dict(batch_first=True))

configs/minkunet/minkunet34_w32_spconv_8xb2-amp-laser-polar-mix-3x_semantickitti.py

@@ -4,6 +4,7 @@
 
 model = dict(
     data_preprocessor=dict(batch_first=True),
-    backbone=dict(sparseconv_backend='spconv'))
+    backbone=dict(sparseconv_backend='spconv'),
+    decode_head=dict(batch_first=True))
 
 optim_wrapper = dict(type='AmpOptimWrapper', loss_scale='dynamic')

configs/minkunet/minkunet34_w32_spconv_8xb2-laser-polar-mix-3x_semantickitti.py

@@ -4,4 +4,5 @@
 
 model = dict(
     data_preprocessor=dict(batch_first=True),
-    backbone=dict(sparseconv_backend='spconv'))
+    backbone=dict(sparseconv_backend='spconv'),
+    decode_head=dict(batch_first=True))

configs/spvcnn/spvcnn_w16_8xb2-amp-15e_semantickitti.py

@@ -1,4 +1,4 @@
-_base_ = ['./spvcnn_w32_8xb2-15e_semantickitti.py']
+_base_ = ['./spvcnn_w32_8xb2-amp-15e_semantickitti.py']
 
 model = dict(
     backbone=dict(

configs/spvcnn/spvcnn_w20_8xb2-amp-15e_semantickitti.py

@@ -1,4 +1,4 @@
-_base_ = ['./spvcnn_w32_8xb2-15e_semantickitti.py']
+_base_ = ['./spvcnn_w32_8xb2-amp-15e_semantickitti.py']
 
 model = dict(
     backbone=dict(

mmdet3d/configs/_base_/models/cylinder3d.py

@@ -1,15 +1,15 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from mmdet3d.models import Cylinder3D
 from mmdet3d.models.backbones import Asymm3DSpconv
 from mmdet3d.models.data_preprocessors import Det3DDataPreprocessor
 from mmdet3d.models.decode_heads.cylinder3d_head import Cylinder3DHead
 from mmdet3d.models.losses import LovaszLoss
+from mmdet3d.models.segmentors import VoxelSegmentor
 from mmdet3d.models.voxel_encoders import SegVFE
 
 grid_shape = [480, 360, 32]
 point_cloud_range = [0, -3.14159265359, -4, 50, 3.14159265359, 2]
 model = dict(
-    type=Cylinder3D,
+    type=VoxelSegmentor,
     data_preprocessor=dict(
         type=Det3DDataPreprocessor,
         voxel=True,

mmdet3d/configs/_base_/models/minkunet.py

@@ -3,10 +3,10 @@
 from mmdet3d.models.data_preprocessors.data_preprocessor import \
     Det3DDataPreprocessor
 from mmdet3d.models.decode_heads.minkunet_head import MinkUNetHead
-from mmdet3d.models.segmentors.minkunet import MinkUNet
+from mmdet3d.models.segmentors import VoxelSegmentor
 
 model = dict(
-    type=MinkUNet,
+    type=VoxelSegmentor,
     data_preprocessor=dict(
         type=Det3DDataPreprocessor,
         voxel=True,

mmdet3d/models/backbones/cylinder3d.py

@@ -13,7 +13,6 @@
 from mmcv.ops import (SparseConv3d, SparseConvTensor, SparseInverseConv3d,
                       SubMConv3d)
 from mmengine.model import BaseModule
-from torch import Tensor
 
 from mmdet3d.registry import MODELS
 from mmdet3d.utils import ConfigType
@@ -457,12 +456,14 @@ def __init__(self,
             indice_key='ddcm',
             norm_cfg=norm_cfg)
 
-    def forward(self, voxel_features: Tensor, coors: Tensor,
-                batch_size: int) -> SparseConvTensor:
+    def forward(self, voxel_dict: dict) -> dict:
         """Forward pass."""
-        coors = coors.int()
-        ret = SparseConvTensor(voxel_features, coors, np.array(self.grid_size),
-                               batch_size)
+        voxel_features = voxel_dict['voxel_feats']
+        voxel_coors = voxel_dict['voxel_coors']
+        voxel_coors = voxel_coors.int()
+        batch_size = voxel_dict['coors'][-1, 0].item() + 1
+        ret = SparseConvTensor(voxel_features, voxel_coors,
+                               np.array(self.grid_size), batch_size)
         ret = self.down_context(ret)
 
         down_skip_list = []
@@ -477,5 +478,6 @@ def forward(self, voxel_features: Tensor, coors: Tensor,
 
         ddcm = self.ddcm(up)
         ddcm.features = torch.cat((ddcm.features, up.features), 1)
+        voxel_dict['voxel_feats'] = ddcm
 
-        return ddcm
+        return voxel_dict

mmdet3d/models/backbones/minkunet_backbone.py

@@ -6,7 +6,7 @@
 import torch
 from mmengine.model import BaseModule
 from mmengine.registry import MODELS
-from torch import Tensor, nn
+from torch import nn
 
 from mmdet3d.models.layers.minkowski_engine_block import (
     IS_MINKOWSKI_ENGINE_AVAILABLE, MinkowskiBasicBlock, MinkowskiBottleneck,
@@ -55,8 +55,8 @@ class MinkUNetBackbone(BaseModule):
         decoder_blocks (List[int]): Number of blocks in each decode layer.
         block_type (str): Type of block in encoder and decoder.
         sparseconv_backend (str): Sparse convolutional backend.
-        init_cfg (dict or :obj:`ConfigDict` or List[dict or :obj:`ConfigDict`]
-            , optional): Initialization config dict.
+        init_cfg (dict or :obj:`ConfigDict` or List[dict or :obj:`ConfigDict`],
+            optional): Initialization config dict.
     """
 
     def __init__(self,
@@ -196,17 +196,17 @@ def __init__(self,
                     [decoder_layer[0],
                      nn.Sequential(*decoder_layer[1:])]))
 
-    def forward(self, voxel_features: Tensor, coors: Tensor) -> Tensor:
+    def forward(self, voxel_dict: dict) -> dict:
         """Forward function.
 
         Args:
-            voxel_features (Tensor): Voxel features in shape (N, C).
-            coors (Tensor): Coordinates in shape (N, 4),
-                the columns in the order of (x_idx, y_idx, z_idx, batch_idx).
+            voxel_dict (dict): Dict containing voxel features.
 
         Returns:
-            Tensor: Backbone features.
+            dict: Backbone features.
         """
+        voxel_features = voxel_dict['voxels']
+        coors = voxel_dict['coors']
         if self.sparseconv_backend == 'torchsparse':
             x = torchsparse.SparseTensor(voxel_features, coors)
         elif self.sparseconv_backend == 'spconv':
@@ -240,6 +240,8 @@ def forward(self, voxel_features: Tensor, coors: Tensor) -> Tensor:
             decoder_outs.append(x)
 
         if self.sparseconv_backend == 'spconv':
-            return decoder_outs[-1].features
+            voxel_dict['voxel_feats'] = decoder_outs[-1].features
         else:
-            return decoder_outs[-1].F
+            voxel_dict['voxel_feats'] = decoder_outs[-1].F
+
+        return voxel_dict

mmdet3d/models/data_preprocessors/data_preprocessor.py

@@ -424,9 +424,9 @@ def voxelize(self, points: List[Tensor],
             voxels = torch.cat(voxels, dim=0)
             coors = torch.cat(coors, dim=0)
         elif self.voxel_type == 'minkunet':
-            voxels, coors = [], []
+            voxels, coors, point2voxel_maps, voxel_inds = [], [], [], []
             voxel_size = points[0].new_tensor(self.voxel_layer.voxel_size)
-            for i, (res, data_sample) in enumerate(zip(points, data_samples)):
+            for i, res in enumerate(points):
                 res_coors = torch.round(res[:, :3] / voxel_size).int()
                 res_coors -= res_coors.min(0)[0]
 
@@ -439,24 +439,22 @@ def voxelize(self, points: List[Tensor],
                         inds = np.random.choice(
                             inds, self.max_voxels, replace=False)
                 inds = torch.from_numpy(inds).cuda()
-                if hasattr(data_sample.gt_pts_seg, 'pts_semantic_mask'):
-                    data_sample.gt_pts_seg.voxel_semantic_mask \
-                        = data_sample.gt_pts_seg.pts_semantic_mask[inds]
                 res_voxel_coors = res_coors[inds]
                 res_voxels = res[inds]
                 if self.batch_first:
                     res_voxel_coors = F.pad(
                         res_voxel_coors, (1, 0), mode='constant', value=i)
-                    data_sample.batch_idx = res_voxel_coors[:, 0]
                 else:
                     res_voxel_coors = F.pad(
                         res_voxel_coors, (0, 1), mode='constant', value=i)
-                    data_sample.batch_idx = res_voxel_coors[:, -1]
-                data_sample.point2voxel_map = point2voxel_map.long()
                 voxels.append(res_voxels)
                 coors.append(res_voxel_coors)
+                point2voxel_maps.append(point2voxel_map)
+                voxel_inds.append(inds)
             voxels = torch.cat(voxels, dim=0)
             coors = torch.cat(coors, dim=0)
+            voxel_dict['point2voxel_maps'] = point2voxel_maps
+            voxel_dict['voxel_inds'] = voxel_inds
 
         else:
             raise ValueError(f'Invalid voxelization type {self.voxel_type}')

mmdet3d/models/decode_heads/cylinder3d_head.py

@@ -9,7 +9,7 @@
 from mmdet3d.models.data_preprocessors.voxelize import dynamic_scatter_3d
 from mmdet3d.registry import MODELS
 from mmdet3d.structures.det3d_data_sample import SampleList
-from mmdet3d.utils import OptConfigType
+from mmdet3d.utils import ConfigType, OptConfigType
 from .decode_head import Base3DDecodeHead
 
 
@@ -88,36 +88,32 @@ def loss_by_feat(self, voxel_dict: dict,
 
         return loss
 
-    def predict(
-        self,
-        voxel_dict: dict,
-        batch_data_samples: SampleList,
-    ) -> List[Tensor]:
+    def predict(self, voxel_dict: dict, batch_input_metas: List[dict],
+                test_cfg: ConfigType) -> List[Tensor]:
         """Forward function for testing.
 
         Args:
             voxel_dict (dict): Features from backbone.
-            batch_data_samples (List[:obj:`Det3DDataSample`]): The det3d data
-                samples. It usually includes information such as `metainfo` and
-                `gt_pts_seg`. We use `point2voxel_map` in this function.
+            batch_input_metas (List[dict]): Meta information of a batch of
+                samples.
+            test_cfg (dict or :obj:`ConfigDict`): The testing config.
 
         Returns:
             List[Tensor]: List of point-wise segmentation logits.
         """
         voxel_dict = self.forward(voxel_dict)
-        seg_pred_list = self.predict_by_feat(voxel_dict, batch_data_samples)
+        seg_pred_list = self.predict_by_feat(voxel_dict, batch_input_metas)
         return seg_pred_list
 
     def predict_by_feat(self, voxel_dict: dict,
-                        batch_data_samples: SampleList) -> List[Tensor]:
+                        batch_input_metas: List[dict]) -> Tensor:
         """Predict function.
 
         Args:
             voxel_dict (dict): The dict may contain `logits`,
                 `point2voxel_map`.
-            batch_data_samples (List[:obj:`Det3DDataSample`]): The det3d data
-                samples. It usually includes information such as `metainfo` and
-                `gt_pts_seg`.
+            batch_input_metas (List[dict]): Meta information of a batch of
+                samples.
 
         Returns:
             List[Tensor]: List of point-wise segmentation logits.
@@ -126,7 +122,7 @@ def predict_by_feat(self, voxel_dict: dict,
 
         seg_pred_list = []
         coors = voxel_dict['voxel_coors']
-        for batch_idx in range(len(batch_data_samples)):
+        for batch_idx in range(len(batch_input_metas)):
             batch_mask = coors[:, 0] == batch_idx
             seg_logits_sample = seg_logits[batch_mask]
             point2voxel_map = voxel_dict['point2voxel_maps'][batch_idx].long()