voxel rcnn release

Author: djiajunustc

LICENSE

@@ -36,11 +36,15 @@ This is the official implementation of [**Voxel R-CNN: Towards High Performance
     ├── pcdet
     ├── tools
     ```
+    Generate the data infos by running the following command:
+    ```
+    python -m pcdet.datasets.kitti.kitti_dataset create_kitti_infos tools/cfgs/dataset_configs/kitti_dataset.yaml
+    ```
 
 3. Setup.
 
     ```
-    python setup.py build develop
+    python setup.py develop
     ```
 
 ### Getting Started

README.md

@@ -63,3 +63,28 @@ def multi_classes_nms(cls_scores, box_preds, nms_config, score_thresh=None):
     pred_boxes = torch.cat(pred_boxes, dim=0)
 
     return pred_scores, pred_labels, pred_boxes
+
+
+def fast_bev_nms(box_scores, box_preds, nms_config, score_thresh=None):
+    src_box_scores = box_scores
+    if score_thresh is not None:
+        scores_mask = (box_scores >= score_thresh)
+        box_scores = box_scores[scores_mask]
+        box_preds = box_preds[scores_mask]
+
+    selected = []
+    if box_scores.shape[0] > 0:
+        box_scores_nms, indices = torch.topk(box_scores, k=min(nms_config.NMS_PRE_MAXSIZE, box_scores.shape[0]))
+        boxes_for_nms = box_preds[indices]
+        bev_iou = iou3d_nms_utils.boxes_iou_bev(boxes_for_nms, boxes_for_nms)
+        bev_iou.triu_(diagonal=1)
+        iou_max, _ = torch.max(bev_iou, dim=0)
+        selected = indices[iou_max <= nms_config.NMS_THRESH]
+        selected = selected[:nms_config.NMS_POST_MAXSIZE]
+        # import pdb
+        # pdb.set_trace()
+
+    if score_thresh is not None:
+        original_idxs = scores_mask.nonzero(as_tuple=False)[:, 0].view(-1)
+        selected = original_idxs[selected]
+    return selected, src_box_scores[selected]

pcdet/models/model_utils/model_nms_utils.py

@@ -4,7 +4,7 @@
 import torch.nn.functional as F
 
 from ...utils import box_coder_utils, common_utils, loss_utils
-from ..model_utils.model_nms_utils import class_agnostic_nms
+from ..model_utils.model_nms_utils import class_agnostic_nms, fast_bev_nms
 from .target_assigner.proposal_target_layer import ProposalTargetLayer
 
 
@@ -63,7 +63,7 @@ def proposal_layer(self, batch_dict, nms_config):
         """
         batch_size = batch_dict['batch_size']
         batch_box_preds = batch_dict['batch_box_preds']
-        batch_cls_preds = batch_dict['batch_cls_preds']
+        batch_cls_preds = torch.sigmoid(batch_dict['batch_cls_preds'])
         rois = batch_box_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE, batch_box_preds.shape[-1]))
         roi_scores = batch_box_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE))
         roi_labels = batch_box_preds.new_zeros((batch_size, nms_config.NMS_POST_MAXSIZE), dtype=torch.long)
@@ -83,9 +83,22 @@ def proposal_layer(self, batch_dict, nms_config):
             if nms_config.MULTI_CLASSES_NMS:
                 raise NotImplementedError
             else:
-                selected, selected_scores = class_agnostic_nms(
-                    box_scores=cur_roi_scores, box_preds=box_preds, nms_config=nms_config
-                )
+                if self.training:
+                    selected, selected_scores = class_agnostic_nms(
+                        box_scores=cur_roi_scores, box_preds=box_preds, nms_config=nms_config
+                    )
+                else:
+                    if nms_config.get("USE_FAST_NMS", False):
+                        selected, selected_scores = fast_bev_nms(
+                            box_scores=cur_roi_scores, box_preds=box_preds, nms_config=nms_config, score_thresh=nms_config.SCORE_THRESH
+                        )
+                    else:
+                        selected, selected_scores = class_agnostic_nms(
+                            box_scores=cur_roi_scores, box_preds=box_preds, nms_config=nms_config
+                        )
+                # selected, selected_scores = class_agnostic_nms(
+                #     box_scores=cur_roi_scores, box_preds=box_preds, nms_config=nms_config
+                # )
 
             rois[index, :len(selected), :] = box_preds[selected]
             roi_scores[index, :len(selected)] = cur_roi_scores[selected]
@@ -189,36 +202,6 @@ def get_box_reg_layer_loss(self, forward_ret_dict):
 
                 rcnn_loss_reg += loss_corner
                 tb_dict['rcnn_loss_corner'] = loss_corner.item()
-
-            if loss_cfgs.GRID_3D_IOU_LOSS and fg_sum > 0:
-                fg_rcnn_reg = rcnn_reg.view(rcnn_batch_size, -1)[fg_mask]
-                fg_roi_boxes3d = roi_boxes3d.view(-1, code_size)[fg_mask]
-
-                fg_roi_boxes3d = fg_roi_boxes3d.view(1, -1, code_size)
-                batch_anchors = fg_roi_boxes3d.clone().detach()
-                roi_ry = fg_roi_boxes3d[:, :, 6].view(-1)
-                roi_xyz = fg_roi_boxes3d[:, :, 0:3].view(-1, 3)
-                batch_anchors[:, :, 0:3] = 0
-                rcnn_boxes3d = self.box_coder.decode_torch(
-                    fg_rcnn_reg.view(batch_anchors.shape[0], -1, code_size), batch_anchors
-                ).view(-1, code_size)
-
-                rcnn_boxes3d = common_utils.rotate_points_along_z(
-                    rcnn_boxes3d.unsqueeze(dim=1), roi_ry
-                ).squeeze(dim=1)
-                rcnn_boxes3d[:, 0:3] += roi_xyz
-
-                loss_iou3d = loss_utils.get_gridify_iou3d_loss(
-                    gt_of_rois_src[fg_mask][:, :7],
-                    rcnn_boxes3d[:, :7]
-                    
-                )
-
-                loss_iou3d = loss_iou3d.mean()
-                loss_iou3d = loss_iou3d * loss_cfgs.LOSS_WEIGHTS['rcnn_iou3d_weight']
-
-                rcnn_loss_reg += loss_iou3d
-                tb_dict['rcnn_loss_iou3d'] = loss_iou3d.item()
         else:
             raise NotImplementedError
 

pcdet/models/roi_heads/roi_head_template.py

@@ -1 +1 @@
-__version__ = "0.1.0+c41ee6c"
+__version__ = "0.1.0+19ea2bd"

pcdet/version.py

@@ -128,7 +128,9 @@ MODEL:
             TEST:
                 NMS_TYPE: nms_gpu
                 MULTI_CLASSES_NMS: False
-                NMS_PRE_MAXSIZE: 4096
+                USE_FAST_NMS: True
+                SCORE_THRESH: 0.0
+                NMS_PRE_MAXSIZE: 2048
                 NMS_POST_MAXSIZE: 100
                 NMS_THRESH: 0.7
 
@@ -139,8 +141,8 @@ MODEL:
             POOL_LAYERS:
                 x_conv3:
                     MLPS: [[32, 32], [32, 32]]
-                    QUERY_RANGES: [[3, 3, 3], [6, 6, 6]]
-                    POOL_RADIUS: [0.6, 1.2]
+                    QUERY_RANGES: [[2, 2, 2], [4, 4, 4]]
+                    POOL_RADIUS: [0.4, 0.8]
                     NSAMPLE: [16, 16]
                     POOL_METHOD: max_pool
                 x_conv4:
@@ -212,4 +214,4 @@ OPTIMIZATION:
     LR_WARMUP: False
     WARMUP_EPOCH: 1
 
-    GRAD_NORM_CLIP: 10
+    GRAD_NORM_CLIP: 10

requirements.txt

@@ -108,7 +108,9 @@ MODEL:
             TEST:
                 NMS_TYPE: nms_gpu
                 MULTI_CLASSES_NMS: False
-                NMS_PRE_MAXSIZE: 4096
+                USE_FAST_NMS: True
+                SCORE_THRESH: 0.0
+                NMS_PRE_MAXSIZE: 2048
                 NMS_POST_MAXSIZE: 100
                 NMS_THRESH: 0.7
 
@@ -119,8 +121,8 @@ MODEL:
             POOL_LAYERS:
                 x_conv3:
                     MLPS: [[32, 32], [32, 32]]
-                    QUERY_RANGES: [[3, 3, 3], [6, 6, 6]]
-                    POOL_RADIUS: [0.6, 1.2]
+                    QUERY_RANGES: [[2, 2, 2], [4, 4, 4]]
+                    POOL_RADIUS: [0.4, 0.8]
                     NSAMPLE: [16, 16]
                     POOL_METHOD: max_pool
                 x_conv4: