[Fix]: Fix PonitRCNN GPU memory occupancy problem (open-mmlab#1928)

* fix pointrcnn oom * fix * fix commenst * fix parta2 ut
ZwwWayne · Dec 16, 2022 · af58208 · af58208
1 parent 39d594f
commit af58208
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 7 deletions.
diff --git a/docs/en/changelog.md b/docs/en/changelog.md
@@ -20,9 +20,9 @@
 
 #### Contributors
 
-A total of 9 developers contributed to this release.
+A total of 11 developers contributed to this release.
 
-@ZwwWayne, @Tai-Wang, @filaPro, @VVsssssk, @ZCMax, @Xiangxu-0103, @holtvogt, @tpoisonooo, @lianqing01
+@ZwwWayne, @Tai-Wang, @filaPro, @VVsssssk, @ZCMax, @Xiangxu-0103, @holtvogt, @tpoisonooo, @lianqing01, @TommyZihao, @aditya9710
 
 ### v1.0.0rc4 (8/8/2022)
 

diff --git a/mmdet3d/models/roi_heads/bbox_heads/point_rcnn_bbox_head.py b/mmdet3d/models/roi_heads/bbox_heads/point_rcnn_bbox_head.py
@@ -203,7 +203,7 @@ def init_weights(self):
  nn.init.constant_(m.bias, 0)
  normal_init(self.conv_reg.weight, mean=0, std=0.001)
 
- def forward(self, feats):
+ def forward(self, input_data):
  """Forward pass.
 
  Args:
@@ -212,9 +212,8 @@ def forward(self, feats):
  Returns:
  tuple[torch.Tensor]: Score of class and bbox predictions.
  """
- input_data = feats.clone().detach()
  xyz_input = input_data[..., 0:self.in_channels].transpose(
- 1, 2).unsqueeze(dim=3).contiguous().clone().detach()
+ 1, 2).unsqueeze(dim=3).contiguous()
  xyz_features = self.xyz_up_layer(xyz_input)
  rpn_features = input_data[..., self.in_channels:].transpose(
  1, 2).unsqueeze(dim=3)

diff --git a/mmdet3d/models/roi_heads/point_rcnn_roi_head.py b/mmdet3d/models/roi_heads/point_rcnn_roi_head.py
@@ -95,7 +95,6 @@ def forward_train(self, feats_dict, input_metas, proposal_list,
  point_cls_preds = feats_dict['points_cls_preds']
  sem_scores = point_cls_preds.sigmoid()
  point_scores = sem_scores.max(-1)[0]
-
  sample_results = self._assign_and_sample(proposal_list, gt_bboxes_3d,
  gt_labels_3d)
 
@@ -211,6 +210,7 @@ def _bbox_forward(self, features, points, batch_size, rois):
  bbox_results = dict(cls_score=cls_score, bbox_pred=bbox_pred)
  return bbox_results
 
+ @torch.no_grad()
  def _assign_and_sample(self, proposal_list, gt_bboxes_3d, gt_labels_3d):
  """Assign and sample proposals for training.
 

diff --git a/tests/test_models/test_heads/test_parta2_bbox_head.py b/tests/test_models/test_heads/test_parta2_bbox_head.py
@@ -2,7 +2,14 @@
 import pytest
 import torch
 from mmcv import Config
-from mmcv.ops import SubMConv3d
+
+from mmdet3d.ops.spconv import IS_SPCONV2_AVAILABLE
+
+if IS_SPCONV2_AVAILABLE:
+ from spconv.pytorch import SubMConv3d
+else:
+ from mmcv.ops import SubMConv3d
+
 from torch.nn import BatchNorm1d, ReLU
 
 from mmdet3d.core.bbox import Box3DMode, LiDARInstance3DBoxes