Fix minor bugs of computing iou3d

mmdet3d/ops/iou3d/iou3d_utils.py

@@ -4,15 +4,17 @@
 
 
 def boxes_iou_bev(boxes_a, boxes_b):
-    """
-    :param boxes_a: (M, 5)
-    :param boxes_b: (N, 5)
-    :return:
-        ans_iou: (M, N)
-    """
+    """Calculate boxes IoU in the bird view.
 
-    ans_iou = torch.cuda.FloatTensor(
-        torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_()
+    Args:
+        boxes_a (torch.Tensor): Input boxes a with shape (M, 5).
+        boxes_b (torch.Tensor): Input boxes b with shape (N, 5).
+
+    Returns:
+        ans_iou (torch.Tensor): IoU result with shape (M, N).
+    """
+    ans_iou = boxes_a.new_zeros(
+        torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
 
     iou3d_cuda.boxes_iou_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(),
                                  ans_iou)
@@ -21,34 +23,41 @@ def boxes_iou_bev(boxes_a, boxes_b):
 
 
 def nms_gpu(boxes, scores, thresh):
+    """Non maximum suppression on GPU.
+
+    Args:
+        boxes (torch.Tensor): Input boxes with shape (N, 5).
+        scores (torch.Tensor): Scores of predicted boxes with shape (N).
+        thresh (torch.Tensor): Threshold of non maximum suppression.
+
+    Returns:
+        torch.Tensor: Remaining indices with scores in descending order.
     """
-    :param boxes: (N, 5) [x1, y1, x2, y2, ry]
-    :param scores: (N)
-    :param thresh:
-    :return:
-    """
-    # areas = (x2 - x1) * (y2 - y1)
     order = scores.sort(0, descending=True)[1]
 
     boxes = boxes[order].contiguous()
 
-    keep = torch.LongTensor(boxes.size(0))
-    num_out = iou3d_cuda.nms_gpu(boxes, keep, thresh)
-    return order[keep[:num_out].cuda()].contiguous()
+    keep = boxes.new_zeros(boxes.size(0))
+    num_out = iou3d_cuda.nms_gpu(boxes, keep, thresh, boxes.device.index)
+    return order[keep[:num_out].cuda(boxes.device)].contiguous()
 
 
 def nms_normal_gpu(boxes, scores, thresh):
+    """Normal non maximum suppression on GPU.
+
+    Args:
+        boxes (torch.Tensor): Input boxes with shape (N, 5).
+        scores (torch.Tensor): Scores of predicted boxes with shape (N).
+        thresh (torch.Tensor): Threshold of non maximum suppression.
+
+    Returns:
+        torch.Tensor: Remaining indices with scores in descending order.
     """
-    :param boxes: (N, 5) [x1, y1, x2, y2, ry]
-    :param scores: (N)
-    :param thresh:
-    :return:
-    """
-    # areas = (x2 - x1) * (y2 - y1)
     order = scores.sort(0, descending=True)[1]
 
     boxes = boxes[order].contiguous()
 
-    keep = torch.LongTensor(boxes.size(0))
-    num_out = iou3d_cuda.nms_normal_gpu(boxes, keep, thresh)
-    return order[keep[:num_out].cuda()].contiguous()
+    keep = boxes.new_zeros(boxes.size(0))
+    num_out = iou3d_cuda.nms_normal_gpu(boxes, keep, thresh,
+                                        boxes.device.index)
+    return order[keep[:num_out].cuda(boxes.device)].contiguous()

mmdet3d/ops/iou3d/src/iou3d.cpp

@@ -92,12 +92,14 @@ int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b,
   return 1;
 }
 
-int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh) {
+int nms_gpu(at::Tensor boxes, at::Tensor keep,
+	    float nms_overlap_thresh, int device_id) {
   // params boxes: (N, 5) [x1, y1, x2, y2, ry]
   // params keep: (N)
 
   CHECK_INPUT(boxes);
   CHECK_CONTIGUOUS(keep);
+  cudaSetDevice(device_id);
 
   int boxes_num = boxes.size(0);
   const float *boxes_data = boxes.data_ptr<float>();
@@ -145,12 +147,13 @@ int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh) {
 }
 
 int nms_normal_gpu(at::Tensor boxes, at::Tensor keep,
-                   float nms_overlap_thresh) {
+                   float nms_overlap_thresh, int device_id) {
   // params boxes: (N, 5) [x1, y1, x2, y2, ry]
   // params keep: (N)
 
   CHECK_INPUT(boxes);
   CHECK_CONTIGUOUS(keep);
+  cudaSetDevice(device_id);
 
   int boxes_num = boxes.size(0);
   const float *boxes_data = boxes.data_ptr<float>();