add group_points, iou3d, roiaware_pool3d and voxelize in parrots (#1504)

mmcv/ops/csrc/pytorch/voxelization.cpp

@@ -47,42 +47,57 @@ void dynamic_voxelize_forward_cpu(const at::Tensor &points, at::Tensor &coors,
                                   const std::vector<float> coors_range,
                                   const int NDim = 3);
 
-int hard_voxelize_forward(const at::Tensor &points, at::Tensor &voxels,
-                          at::Tensor &coors, at::Tensor &num_points_per_voxel,
-                          const std::vector<float> voxel_size,
-                          const std::vector<float> coors_range,
-                          const int max_points, const int max_voxels,
-                          const int NDim = 3) {
+void hard_voxelize_forward(const at::Tensor &points,
+                           const at::Tensor &voxel_size,
+                           const at::Tensor &coors_range, at::Tensor &voxels,
+                           at::Tensor &coors, at::Tensor &num_points_per_voxel,
+                           at::Tensor &voxel_num, const int max_points,
+                           const int max_voxels, const int NDim = 3) {
+  int64_t *voxel_num_data = voxel_num.data_ptr<int64_t>();
+  std::vector<float> voxel_size_v(
+      voxel_size.data_ptr<float>(),
+      voxel_size.data_ptr<float>() + voxel_size.numel());
+  std::vector<float> coors_range_v(
+      coors_range.data_ptr<float>(),
+      coors_range.data_ptr<float>() + coors_range.numel());
   if (points.device().is_cuda()) {
 #ifdef MMCV_WITH_CUDA
     CHECK_CUDA_INPUT(points);
 
-    return hard_voxelize_forward_cuda(
-        points, voxels, coors, num_points_per_voxel, voxel_size, coors_range,
-        max_points, max_voxels, NDim);
+    *voxel_num_data = hard_voxelize_forward_cuda(
+        points, voxels, coors, num_points_per_voxel, voxel_size_v,
+        coors_range_v, max_points, max_voxels, NDim);
 #else
     AT_ERROR("hard_voxelize is not compiled with GPU support");
 #endif
   } else {
-    return hard_voxelize_forward_cpu(points, voxels, coors,
-                                     num_points_per_voxel, voxel_size,
-                                     coors_range, max_points, max_voxels, NDim);
+    *voxel_num_data = hard_voxelize_forward_cpu(
+        points, voxels, coors, num_points_per_voxel, voxel_size_v,
+        coors_range_v, max_points, max_voxels, NDim);
   }
 }
 
-void dynamic_voxelize_forward(const at::Tensor &points, at::Tensor &coors,
-                              const std::vector<float> voxel_size,
-                              const std::vector<float> coors_range,
+void dynamic_voxelize_forward(const at::Tensor &points,
+                              const at::Tensor &voxel_size,
+                              const at::Tensor &coors_range, at::Tensor &coors,
                               const int NDim = 3) {
+  std::vector<float> voxel_size_v(
+      voxel_size.data_ptr<float>(),
+      voxel_size.data_ptr<float>() + voxel_size.numel());
+  std::vector<float> coors_range_v(
+      coors_range.data_ptr<float>(),
+      coors_range.data_ptr<float>() + coors_range.numel());
   if (points.device().is_cuda()) {
 #ifdef MMCV_WITH_CUDA
     CHECK_CUDA_INPUT(points);
 
-    dynamic_voxelize_forward_cuda(points, coors, voxel_size, coors_range, NDim);
+    dynamic_voxelize_forward_cuda(points, coors, voxel_size_v, coors_range_v,
+                                  NDim);
 #else
     AT_ERROR("dynamic_voxelize is not compiled with GPU support");
 #endif
   } else {
-    dynamic_voxelize_forward_cpu(points, coors, voxel_size, coors_range, NDim);
+    dynamic_voxelize_forward_cpu(points, coors, voxel_size_v, coors_range_v,
+                                 NDim);
   }
 }

mmcv/ops/group_points.py

@@ -192,8 +192,15 @@ class GroupingOperation(Function):
         _, C, N = features.size()
         output = torch.cuda.FloatTensor(B, C, nfeatures, nsample)
 
-        ext_module.group_points_forward(B, C, N, nfeatures, nsample, features,
-                                        indices, output)
+        ext_module.group_points_forward(
+            features,
+            indices,
+            output,
+            b=B,
+            c=C,
+            n=N,
+            npoints=nfeatures,
+            nsample=nsample)
 
         ctx.for_backwards = (indices, N)
         return output
@@ -215,9 +222,15 @@ class GroupingOperation(Function):
         grad_features = torch.cuda.FloatTensor(B, C, N).zero_()
 
         grad_out_data = grad_out.data.contiguous()
-        ext_module.group_points_backward(B, C, N, npoint, nsample,
-                                         grad_out_data, idx,
-                                         grad_features.data)
+        ext_module.group_points_backward(
+            grad_out_data,
+            idx,
+            grad_features.data,
+            b=B,
+            c=C,
+            n=N,
+            npoints=npoint,
+            nsample=nsample)
         return grad_features, None
 
 

mmcv/ops/iou3d.py

@@ -55,7 +55,9 @@ def nms_bev(boxes, scores, thresh, pre_max_size=None, post_max_size=None):
     boxes = boxes[order].contiguous()
 
     keep = torch.zeros(boxes.size(0), dtype=torch.long)
-    num_out = ext_module.iou3d_nms_forward(boxes, keep, thresh)
+    num_out = torch.zeros(size=(), dtype=torch.long)
+    ext_module.iou3d_nms_forward(
+        boxes, keep, num_out, nms_overlap_thresh=thresh)
     keep = order[keep[:num_out].cuda(boxes.device)].contiguous()
     if post_max_size is not None:
         keep = keep[:post_max_size]
@@ -81,5 +83,7 @@ def nms_normal_bev(boxes, scores, thresh):
     boxes = boxes[order].contiguous()
 
     keep = torch.zeros(boxes.size(0), dtype=torch.long)
-    num_out = ext_module.iou3d_nms_normal_forward(boxes, keep, thresh)
+    num_out = torch.zeros(size=(), dtype=torch.long)
+    ext_module.iou3d_nms_normal_forward(
+        boxes, keep, num_out, nms_overlap_thresh=thresh)
     return order[keep[:num_out].cuda(boxes.device)].contiguous()

mmcv/ops/roiaware_pool3d.py

@@ -93,9 +93,14 @@ class RoIAwarePool3dFunction(Function):
             (num_rois, out_x, out_y, out_z, max_pts_per_voxel),
             dtype=torch.int)
 
-        ext_module.roiaware_pool3d_forward(rois, pts, pts_feature, argmax,
-                                           pts_idx_of_voxels, pooled_features,
-                                           mode)
+        ext_module.roiaware_pool3d_forward(
+            rois,
+            pts,
+            pts_feature,
+            argmax,
+            pts_idx_of_voxels,
+            pooled_features,
+            pool_method=mode)
 
         ctx.roiaware_pool3d_for_backward = (pts_idx_of_voxels, argmax, mode,
                                             num_pts, num_channels)
@@ -107,8 +112,11 @@ class RoIAwarePool3dFunction(Function):
         pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret
 
         grad_in = grad_out.new_zeros((num_pts, num_channels))
-        ext_module.roiaware_pool3d_backward(pts_idx_of_voxels, argmax,
-                                            grad_out.contiguous(), grad_in,
-                                            mode)
+        ext_module.roiaware_pool3d_backward(
+            pts_idx_of_voxels,
+            argmax,
+            grad_out.contiguous(),
+            grad_in,
+            pool_method=mode)
 
         return None, None, grad_in, None, None, None

mmcv/ops/voxelize.py

@@ -46,8 +46,12 @@ class _Voxelization(Function):
         """
         if max_points == -1 or max_voxels == -1:
             coors = points.new_zeros(size=(points.size(0), 3), dtype=torch.int)
-            ext_module.dynamic_voxelize_forward(points, coors, voxel_size,
-                                                coors_range, 3)
+            ext_module.dynamic_voxelize_forward(
+                points,
+                torch.tensor(voxel_size, dtype=torch.float),
+                torch.tensor(coors_range, dtype=torch.float),
+                coors,
+                NDim=3)
             return coors
         else:
             voxels = points.new_zeros(
@@ -55,9 +59,18 @@ class _Voxelization(Function):
             coors = points.new_zeros(size=(max_voxels, 3), dtype=torch.int)
             num_points_per_voxel = points.new_zeros(
                 size=(max_voxels, ), dtype=torch.int)
-            voxel_num = ext_module.hard_voxelize_forward(
-                points, voxels, coors, num_points_per_voxel, voxel_size,
-                coors_range, max_points, max_voxels, 3)
+            voxel_num = torch.zeros(size=(), dtype=torch.long)
+            ext_module.hard_voxelize_forward(
+                points,
+                torch.tensor(voxel_size, dtype=torch.float),
+                torch.tensor(coors_range, dtype=torch.float),
+                voxels,
+                coors,
+                num_points_per_voxel,
+                voxel_num,
+                max_points=max_points,
+                max_voxels=max_voxels,
+                NDim=3)
             # select the valid voxels
             voxels_out = voxels[:voxel_num]
             coors_out = coors[:voxel_num]