add setup.py

requirements.txt

@@ -1,7 +1,7 @@
 munch
 pandas
 plyfile
-pyyaml==5.4.1
+pyyaml
 scikit-learn
 scipy
 six

setup.py

@@ -0,0 +1,26 @@
+from setuptools import setup
+
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+if __name__ == '__main__':
+    setup(
+        name='softgroup',
+        version='1.0',
+        description='SoftGroup: SoftGroup for 3D Instance Segmentation [CVPR 2022]',
+        author='Thang Vu',
+        author_email='thangvubk@kaist.ac.kr',
+        packages=['softgroup'],
+        package_data={'softgroup.ops': ['*/*.so']},
+        ext_modules=[
+            CUDAExtension(
+                name='softgroup.ops.ops',
+                sources=[
+                    'softgroup/ops/src/softgroup_api.cpp', 'softgroup/ops/src/softgroup_ops.cpp',
+                    'softgroup/ops/src/cuda.cu'
+                ],
+                extra_compile_args={
+                    'cxx': ['-g'],
+                    'nvcc': ['-O2']
+                })
+        ],
+        cmdclass={'build_ext': BuildExtension})

softgroup/data/custom.py

@@ -8,7 +8,7 @@ import scipy.ndimage
 import torch
 from torch.utils.data import Dataset
 
-from ..lib.softgroup_ops import voxelization_idx
+from ..ops import voxelization_idx
 
 
 class CustomDataset(Dataset):

softgroup/data/s3dis.py

@@ -4,7 +4,7 @@ from glob import glob
 import numpy as np
 import torch
 
-from ..lib.softgroup_ops import voxelization_idx
+from ..ops import voxelization_idx
 from .custom import CustomDataset
 
 

softgroup/lib/softgroup_ops/__init__.py

@@ -1 +0,0 @@
-from .functions.softgroup_ops import *

softgroup/lib/softgroup_ops/src/get_iou/get_iou.cpp

@@ -1,23 +0,0 @@
-/*
-Get the IoU between predictions and gt masks
-Written by Li Jiang
-All Rights Reserved 2020.
-*/
-
-#include "get_iou.h"
-
-void get_iou(at::Tensor proposals_idx_tensor,
-             at::Tensor proposals_offset_tensor,
-             at::Tensor instance_labels_tensor,
-             at::Tensor instance_pointnum_tensor,
-             at::Tensor proposals_iou_tensor, int nInstance, int nProposal) {
-  int *proposals_idx = proposals_idx_tensor.data<int>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
-  long *instance_labels = instance_labels_tensor.data<long>();
-  int *instance_pointnum = instance_pointnum_tensor.data<int>();
-
-  float *proposals_iou = proposals_iou_tensor.data<float>();
-
-  get_iou_cuda(nInstance, nProposal, proposals_idx, proposals_offset,
-               instance_labels, instance_pointnum, proposals_iou);
-}

softgroup/lib/softgroup_ops/src/get_iou/get_iou.cu

@@ -1,48 +0,0 @@
-/*
-Get the IoU between predictions and gt masks
-Written by Li Jiang
-All Rights Reserved 2020.
-*/
-
-#include "get_iou.h"
-#include <math.h>
-#include <stdio.h>
-
-__global__ void get_iou_cuda_(int nInstance, int nProposal, int *proposals_idx,
-                              int *proposals_offset, long *instance_labels,
-                              int *instance_pointnum, float *proposals_iou) {
-  for (int proposal_id = blockIdx.x; proposal_id < nProposal;
-       proposal_id += gridDim.x) {
-    int start = proposals_offset[proposal_id];
-    int end = proposals_offset[proposal_id + 1];
-    int proposal_total = end - start;
-    for (int instance_id = threadIdx.x; instance_id < nInstance;
-         instance_id += blockDim.x) {
-      int instance_total = instance_pointnum[instance_id];
-      int intersection = 0;
-      for (int i = start; i < end; i++) {
-        int idx = proposals_idx[i];
-        if ((int)instance_labels[idx] == instance_id) {
-          intersection += 1;
-        }
-      }
-      proposals_iou[proposal_id * nInstance + instance_id] =
-          (float)intersection /
-          ((float)(proposal_total + instance_total - intersection) + 1e-5);
-    }
-  }
-}
-
-// input: proposals_idx (sumNPoint), int
-// input: proposals_offset (nProposal + 1), int
-// input: instance_labels (N), long, 0~total_nInst-1, -100
-// input: instance_pointnum (total_nInst), int
-// output: proposals_iou (nProposal, total_nInst), float
-void get_iou_cuda(int nInstance, int nProposal, int *proposals_idx,
-                  int *proposals_offset, long *instance_labels,
-                  int *instance_pointnum, float *proposals_iou) {
-  get_iou_cuda_<<<std::min(nProposal, (int)32768),
-                  std::min(nInstance, (int)256)>>>(
-      nInstance, nProposal, proposals_idx, proposals_offset, instance_labels,
-      instance_pointnum, proposals_iou);
-}

softgroup/lib/softgroup_ops/src/get_iou/get_iou.h

@@ -1,24 +0,0 @@
-/*
-Get the IoU between predictions and gt masks
-Written by Li Jiang
-All Rights Reserved 2020.
-*/
-
-#ifndef GET_IOU_H
-#define GET_IOU_H
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/serialize/tensor.h>
-
-#include "../datatype/datatype.h"
-
-//
-void get_iou_cuda(int nInstance, int nProposal, int *proposals_idx,
-                  int *proposals_offset, long *instance_labels,
-                  int *instance_pointnum, float *proposals_iou);
-void get_iou(at::Tensor proposals_idx_tensor,
-             at::Tensor proposals_offset_tensor,
-             at::Tensor instance_labels_tensor,
-             at::Tensor instance_pointnum_tensor,
-             at::Tensor proposals_iou_tensor, int nInstance, int nProposal);
-
-#endif // GET_IOU_H

softgroup/lib/softgroup_ops/src/roipool/roipool.cpp

@@ -1,53 +0,0 @@
-/*
-ROI Max Pool
-Written by Li Jiang
-All Rights Reserved 2020.
-*/
-
-#include "roipool.h"
-
-void roipool_fp(at::Tensor feats_tensor, at::Tensor proposals_offset_tensor,
-                at::Tensor output_feats_tensor, at::Tensor output_maxidx_tensor,
-                int nProposal, int C) {
-  float *feats = feats_tensor.data<float>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
-  float *output_feats = output_feats_tensor.data<float>();
-  int *output_maxidx = output_maxidx_tensor.data<int>();
-
-  roipool_fp_cuda(nProposal, C, feats, proposals_offset, output_feats,
-                  output_maxidx);
-}
-
-void roipool_bp(at::Tensor d_feats_tensor, at::Tensor proposals_offset_tensor,
-                at::Tensor output_maxidx_tensor,
-                at::Tensor d_output_feats_tensor, int nProposal, int C) {
-  float *d_feats = d_feats_tensor.data<float>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
-  int *output_maxidx = output_maxidx_tensor.data<int>();
-  float *d_output_feats = d_output_feats_tensor.data<float>();
-
-  roipool_bp_cuda(nProposal, C, d_feats, proposals_offset, output_maxidx,
-                  d_output_feats);
-}
-
-void global_avg_pool_fp(at::Tensor feats_tensor,
-                        at::Tensor proposals_offset_tensor,
-                        at::Tensor output_feats_tensor, int nProposal, int C) {
-  float *feats = feats_tensor.data<float>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
-  float *output_feats = output_feats_tensor.data<float>();
-
-  global_avg_pool_fp_cuda(nProposal, C, feats, proposals_offset, output_feats);
-}
-
-void global_avg_pool_bp(at::Tensor d_feats_tensor,
-                        at::Tensor proposals_offset_tensor,
-                        at::Tensor d_output_feats_tensor, int nProposal,
-                        int C) {
-  float *d_feats = d_feats_tensor.data<float>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
-  float *d_output_feats = d_output_feats_tensor.data<float>();
-
-  global_avg_pool_bp_cuda(nProposal, C, d_feats, proposals_offset,
-                          d_output_feats);
-}

softgroup/model/softgroup.py

@@ -6,9 +6,9 @@ import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
-from ..lib.softgroup_ops import (ballquery_batch_p, bfs_cluster, get_mask_iou_on_cluster,
+from ..ops import (ballquery_batch_p, bfs_cluster, get_mask_iou_on_cluster, get_mask_iou_on_pred,
-                                 get_mask_iou_on_pred, get_mask_label, global_avg_pool, sec_max,
+                   get_mask_label, global_avg_pool, sec_max, sec_min, voxelization,
-                                 sec_min, voxelization, voxelization_idx)
+                   voxelization_idx)
 from ..util import cuda_cast, force_fp32, rle_encode
 from .blocks import MLP, ResidualBlock, UBlock
 

softgroup/ops/__init__.py

@@ -0,0 +1 @@
+from .functions import *

softgroup/ops/functions.py

@@ -1,7 +1,7 @@
 import torch
 from torch.autograd import Function
 
-from .. import SOFTGROUP_OP
+from . import ops
 
 
 class GetMaskIoUOnCluster(Function):
@@ -30,7 +30,7 @@ class GetMaskIoUOnCluster(Function):
         assert instance_labels.is_contiguous() and instance_labels.is_cuda
         assert instance_pointnum.is_contiguous() and instance_pointnum.is_cuda
 
-        SOFTGROUP_OP.get_mask_iou_on_cluster(proposals_idx, proposals_offset, instance_labels,
+        ops.get_mask_iou_on_cluster(proposals_idx, proposals_offset, instance_labels,
                                     instance_pointnum, proposals_iou, nInstance, nProposal)
 
         return proposals_iou
@@ -71,7 +71,7 @@ class GetMaskIoUOnPred(Function):
         assert instance_pointnum.is_contiguous() and instance_pointnum.is_cuda
         assert mask_scores_sigmoid.is_contiguous() and mask_scores_sigmoid.is_cuda
 
-        SOFTGROUP_OP.get_mask_iou_on_pred(proposals_idx, proposals_offset, instance_labels,
+        ops.get_mask_iou_on_pred(proposals_idx, proposals_offset, instance_labels,
                                  instance_pointnum, proposals_iou, nInstance, nProposal,
                                  mask_scores_sigmoid)
 
@@ -112,7 +112,7 @@ class GetMaskLabel(Function):
         assert instance_labels.is_contiguous() and instance_labels.is_cuda
         assert instance_cls.is_contiguous() and instance_cls.is_cuda
 
-        SOFTGROUP_OP.get_mask_label(proposals_idx, proposals_offset, instance_labels, instance_cls,
+        ops.get_mask_label(proposals_idx, proposals_offset, instance_labels, instance_cls,
                            proposals_iou, nInstance, nProposal, iou_thr, mask_label)
 
         return mask_label
@@ -146,7 +146,7 @@ class Voxelization_Idx(Function):
         input_map = torch.IntTensor(N).zero_()
         output_map = input_map.new()
 
-        SOFTGROUP_OP.voxelize_idx(coords, output_coords, input_map, output_map, batchsize, mode)
+        ops.voxelize_idx(coords, output_coords, input_map, output_map, batchsize, mode)
         return output_coords, input_map, output_map
 
     @staticmethod
@@ -177,7 +177,7 @@ class Voxelization(Function):
 
         ctx.for_backwards = (map_rule, mode, maxActive, N)
 
-        SOFTGROUP_OP.voxelize_fp(feats, output_feats, map_rule, mode, M, maxActive, C)
+        ops.voxelize_fp(feats, output_feats, map_rule, mode, M, maxActive, C)
         return output_feats
 
     @staticmethod
@@ -187,54 +187,13 @@ class Voxelization(Function):
 
         d_feats = torch.cuda.FloatTensor(N, C).zero_()
 
-        SOFTGROUP_OP.voxelize_bp(d_output_feats.contiguous(), d_feats, map_rule, mode, M, maxActive,
+        ops.voxelize_bp(d_output_feats.contiguous(), d_feats, map_rule, mode, M, maxActive, C)
-                                 C)
         return d_feats, None, None
 
 
 voxelization = Voxelization.apply
 
 
-class PointRecover(Function):
-
-    @staticmethod
-    def forward(ctx, feats, map_rule, nPoint):
-        '''
-        :param ctx:
-        :param feats: cuda float M * C
-        :param map_rule: cuda int M * (maxActive + 1)
-        :param nPoint: int
-        :return: output_feats: cuda float N * C
-        '''
-        assert map_rule.is_contiguous()
-        assert feats.is_contiguous()
-        M, C = feats.size()
-        maxActive = map_rule.size(1) - 1
-
-        output_feats = torch.cuda.FloatTensor(nPoint, C).zero_()
-
-        ctx.for_backwards = (map_rule, maxActive, M)
-
-        SOFTGROUP_OP.point_recover_fp(feats, output_feats, map_rule, M, maxActive, C)
-
-        return output_feats
-
-    @staticmethod
-    def backward(ctx, d_output_feats):
-        map_rule, maxActive, M = ctx.for_backwards
-        N, C = d_output_feats.size()
-
-        d_feats = torch.cuda.FloatTensor(M, C).zero_()
-
-        SOFTGROUP_OP.point_recover_bp(d_output_feats.contiguous(), d_feats, map_rule, M, maxActive,
-                                      C)
-
-        return d_feats, None, None
-
-
-point_recover = PointRecover.apply
-
-
 class BallQueryBatchP(Function):
 
     @staticmethod
@@ -259,8 +218,8 @@ class BallQueryBatchP(Function):
         while True:
             idx = torch.cuda.IntTensor(n * meanActive).zero_()
             start_len = torch.cuda.IntTensor(n, 2).zero_()
-            nActive = SOFTGROUP_OP.ballquery_batch_p(coords, batch_idxs, batch_offsets, idx,
+            nActive = ops.ballquery_batch_p(coords, batch_idxs, batch_offsets, idx, start_len, n,
-                                                     start_len, n, meanActive, radius)
+                                            meanActive, radius)
             if nActive <= n * meanActive:
                 break
             meanActive = int(nActive // n + 1)
@@ -296,7 +255,7 @@ class BFSCluster(Function):
         cluster_idxs = ball_query_idxs.new()
         cluster_offsets = ball_query_idxs.new()
 
-        SOFTGROUP_OP.bfs_cluster(cluster_numpoint_mean, ball_query_idxs, start_len, cluster_idxs,
+        ops.bfs_cluster(cluster_numpoint_mean, ball_query_idxs, start_len, cluster_idxs,
                         cluster_offsets, N, threshold, class_id)
 
         return cluster_idxs, cluster_offsets
@@ -309,48 +268,6 @@ class BFSCluster(Function):
 bfs_cluster = BFSCluster.apply
 
 
-class RoiPool(Function):
-
-    @staticmethod
-    def forward(ctx, feats, proposals_offset):
-        '''
-        :param ctx:
-        :param feats: (sumNPoint, C) float
-        :param proposals_offset: (nProposal + 1) int
-        :return: output_feats (nProposal, C) float
-        '''
-        nProposal = proposals_offset.size(0) - 1
-        sumNPoint, C = feats.size()
-
-        assert feats.is_contiguous()
-        assert proposals_offset.is_contiguous()
-
-        output_feats = torch.cuda.FloatTensor(nProposal, C).zero_()
-        output_maxidx = torch.cuda.IntTensor(nProposal, C).zero_()
-
-        SOFTGROUP_OP.roipool_fp(feats, proposals_offset, output_feats, output_maxidx, nProposal, C)
-
-        ctx.for_backwards = (output_maxidx, proposals_offset, sumNPoint)
-
-        return output_feats
-
-    @staticmethod
-    def backward(ctx, d_output_feats):
-        nProposal, C = d_output_feats.size()
-
-        output_maxidx, proposals_offset, sumNPoint = ctx.for_backwards
-
-        d_feats = torch.cuda.FloatTensor(sumNPoint, C).zero_()
-
-        SOFTGROUP_OP.roipool_bp(d_feats, proposals_offset, output_maxidx,
-                                d_output_feats.contiguous(), nProposal, C)
-
-        return d_feats, None
-
-
-roipool = RoiPool.apply
-
-
 class GlobalAvgPool(Function):
 
     @staticmethod
@@ -369,7 +286,7 @@ class GlobalAvgPool(Function):
 
         output_feats = torch.cuda.FloatTensor(nProposal, C).zero_()
 
-        SOFTGROUP_OP.global_avg_pool_fp(feats, proposals_offset, output_feats, nProposal, C)
+        ops.global_avg_pool_fp(feats, proposals_offset, output_feats, nProposal, C)
 
         ctx.for_backwards = (proposals_offset, sumNPoint)
 
@@ -383,8 +300,7 @@ class GlobalAvgPool(Function):
 
         d_feats = torch.cuda.FloatTensor(sumNPoint, C).zero_()
 
-        SOFTGROUP_OP.global_avg_pool_bp(d_feats, proposals_offset, d_output_feats.contiguous(),
+        ops.global_avg_pool_bp(d_feats, proposals_offset, d_output_feats.contiguous(), nProposal, C)
-                                        nProposal, C)
 
         return d_feats, None
 
@@ -392,41 +308,6 @@ class GlobalAvgPool(Function):
 global_avg_pool = GlobalAvgPool.apply
 
 
-class GetIoU(Function):
-
-    @staticmethod
-    def forward(ctx, proposals_idx, proposals_offset, instance_labels, instance_pointnum):
-        '''
-        :param ctx:
-        :param proposals_idx: (sumNPoint), int
-        :param proposals_offset: (nProposal + 1), int
-        :param instance_labels: (N), long, 0~total_nInst-1, -100
-        :param instance_pointnum: (total_nInst), int
-        :return: proposals_iou: (nProposal, total_nInst), float
-        '''
-        nInstance = instance_pointnum.size(0)
-        nProposal = proposals_offset.size(0) - 1
-
-        assert proposals_idx.is_contiguous() and proposals_idx.is_cuda
-        assert proposals_offset.is_contiguous() and proposals_offset.is_cuda
-        assert instance_labels.is_contiguous() and instance_labels.is_cuda
-        assert instance_pointnum.is_contiguous() and instance_pointnum.is_cuda
-
-        proposals_iou = torch.cuda.FloatTensor(nProposal, nInstance).zero_()
-
-        SOFTGROUP_OP.get_iou(proposals_idx, proposals_offset, instance_labels, instance_pointnum,
-                             proposals_iou, nInstance, nProposal)
-
-        return proposals_iou
-
-    @staticmethod
-    def backward(ctx, a=None):
-        return None, None, None, None
-
-
-get_iou = GetIoU.apply
-
-
 class SecMean(Function):
 
     @staticmethod
@@ -445,7 +326,7 @@ class SecMean(Function):
 
         out = torch.cuda.FloatTensor(nProposal, C).zero_()
 
-        SOFTGROUP_OP.sec_mean(inp, offsets, out, nProposal, C)
+        ops.sec_mean(inp, offsets, out, nProposal, C)
 
         return out
 
@@ -475,7 +356,7 @@ class SecMin(Function):
 
         out = torch.cuda.FloatTensor(nProposal, C).zero_()
 
-        SOFTGROUP_OP.sec_min(inp, offsets, out, nProposal, C)
+        ops.sec_min(inp, offsets, out, nProposal, C)
 
         return out
 
@@ -505,7 +386,7 @@ class SecMax(Function):
 
         out = torch.cuda.FloatTensor(nProposal, C).zero_()
 
-        SOFTGROUP_OP.sec_max(inp, offsets, out, nProposal, C)
+        ops.sec_max(inp, offsets, out, nProposal, C)
 
         return out
 

softgroup/ops/src/bfs_cluster/bfs_cluster.cpp

@@ -18,11 +18,11 @@ int ballquery_batch_p(at::Tensor xyz_tensor, at::Tensor batch_idxs_tensor,
                       at::Tensor batch_offsets_tensor, at::Tensor idx_tensor,
                       at::Tensor start_len_tensor, int n, int meanActive,
                       float radius) {
-  const float *xyz = xyz_tensor.data<float>();
+  const float *xyz = xyz_tensor.data_ptr<float>();
-  const int *batch_idxs = batch_idxs_tensor.data<int>();
+  const int *batch_idxs = batch_idxs_tensor.data_ptr<int>();
-  const int *batch_offsets = batch_offsets_tensor.data<int>();
+  const int *batch_offsets = batch_offsets_tensor.data_ptr<int>();
-  int *idx = idx_tensor.data<int>();
+  int *idx = idx_tensor.data_ptr<int>();
-  int *start_len = start_len_tensor.data<int>();
+  int *start_len = start_len_tensor.data_ptr<int>();
 
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   int cumsum = ballquery_batch_p_cuda(n, meanActive, radius, xyz, batch_idxs,
@@ -109,9 +109,9 @@ void bfs_cluster(at::Tensor class_numpoint_mean_tensor,
                  at::Tensor cluster_idxs_tensor,
                  at::Tensor cluster_offsets_tensor, const int N,
                  float threshold, const int class_id) {
-  float *class_numpoint_mean = class_numpoint_mean_tensor.data<float>();
+  float *class_numpoint_mean = class_numpoint_mean_tensor.data_ptr<float>();
-  Int *ball_query_idxs = ball_query_idxs_tensor.data<Int>();
+  Int *ball_query_idxs = ball_query_idxs_tensor.data_ptr<Int>();
-  int *start_len = start_len_tensor.data<int>();
+  int *start_len = start_len_tensor.data_ptr<int>();
   ConnectedComponents CCs;
   int sumNPoint = get_clusters(class_numpoint_mean, ball_query_idxs, start_len,
                                N, threshold, CCs, class_id);
@@ -120,7 +120,7 @@ void bfs_cluster(at::Tensor class_numpoint_mean_tensor,
   cluster_offsets_tensor.resize_({nCluster + 1});
   cluster_idxs_tensor.zero_();
   cluster_offsets_tensor.zero_();
-  int *cluster_idxs = cluster_idxs_tensor.data<int>();
+  int *cluster_idxs = cluster_idxs_tensor.data_ptr<int>();
-  int *cluster_offsets = cluster_offsets_tensor.data<int>();
+  int *cluster_offsets = cluster_offsets_tensor.data_ptr<int>();
   fill_cluster_idxs_(CCs, cluster_idxs, cluster_offsets);
 }

softgroup/ops/src/bfs_cluster/bfs_cluster.h

@@ -7,7 +7,6 @@ All Rights Reserved 2020.
 #ifndef BFS_CLUSTER_H
 #define BFS_CLUSTER_H
 #include <ATen/cuda/CUDAContext.h>
-#include <THC/THC.h>
 #include <torch/serialize/tensor.h>
 
 #include "../datatype/datatype.h"

softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp

@@ -10,11 +10,11 @@ void get_mask_iou_on_cluster(at::Tensor proposals_idx_tensor,
                              at::Tensor instance_pointnum_tensor,
                              at::Tensor proposals_iou_tensor, int nInstance,
                              int nProposal) {
-  int *proposals_idx = proposals_idx_tensor.data<int>();
+  int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
+  int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
-  long *instance_labels = instance_labels_tensor.data<long>();
+  long *instance_labels = instance_labels_tensor.data_ptr<long>();
-  int *instance_pointnum = instance_pointnum_tensor.data<int>();
+  int *instance_pointnum = instance_pointnum_tensor.data_ptr<int>();
-  float *proposals_iou = proposals_iou_tensor.data<float>();
+  float *proposals_iou = proposals_iou_tensor.data_ptr<float>();
 
   // input: nInstance (1,), int
   // input: nProposal (1,), int
@@ -37,12 +37,12 @@ void get_mask_iou_on_pred(at::Tensor proposals_idx_tensor,
                           at::Tensor proposals_iou_tensor, int nInstance,
                           int nProposal,
                           at::Tensor mask_scores_sigmoid_tensor) {
-  int *proposals_idx = proposals_idx_tensor.data<int>();
+  int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
+  int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
-  long *instance_labels = instance_labels_tensor.data<long>();
+  long *instance_labels = instance_labels_tensor.data_ptr<long>();
-  int *instance_pointnum = instance_pointnum_tensor.data<int>();
+  int *instance_pointnum = instance_pointnum_tensor.data_ptr<int>();
-  float *proposals_iou = proposals_iou_tensor.data<float>();
+  float *proposals_iou = proposals_iou_tensor.data_ptr<float>();
-  float *mask_scores_sigmoid = mask_scores_sigmoid_tensor.data<float>();
+  float *mask_scores_sigmoid = mask_scores_sigmoid_tensor.data_ptr<float>();
 
   // input: nInstance (1,), int
   // input: nProposal (1,), int
@@ -65,12 +65,12 @@ void get_mask_label(at::Tensor proposals_idx_tensor,
                     at::Tensor proposals_iou_tensor, int nInstance,
                     int nProposal, float iou_thr,
                     at::Tensor mask_labels_tensor) {
-  int *proposals_idx = proposals_idx_tensor.data<int>();
+  int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
-  int *proposals_offset = proposals_offset_tensor.data<int>();
+  int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
-  long *instance_labels = instance_labels_tensor.data<long>();
+  long *instance_labels = instance_labels_tensor.data_ptr<long>();
-  long *instance_cls = instance_cls_tensor.data<long>();
+  long *instance_cls = instance_cls_tensor.data_ptr<long>();
-  float *proposals_iou = proposals_iou_tensor.data<float>();
+  float *proposals_iou = proposals_iou_tensor.data_ptr<float>();
-  float *mask_label = mask_labels_tensor.data<float>();
+  float *mask_label = mask_labels_tensor.data_ptr<float>();
 
   // input: nInstance (1,), int
   // input: nProposal (1,), int

softgroup/ops/src/cuda.cu

@@ -3,7 +3,6 @@
 
 #include "bfs_cluster/bfs_cluster.cu"
 #include "cal_iou_and_masklabel/cal_iou_and_masklabel.cu"
-#include "get_iou/get_iou.cu"
 #include "roipool/roipool.cu"
 #include "sec_mean/sec_mean.cu"
 #include "voxelize/voxelize.cu"

softgroup/ops/src/roipool/roipool.cpp

@@ -0,0 +1,29 @@
+/*
+ROI Max Pool
+Written by Li Jiang
+All Rights Reserved 2020.
+*/
+
+#include "roipool.h"
+
+void global_avg_pool_fp(at::Tensor feats_tensor,
+                        at::Tensor proposals_offset_tensor,
+                        at::Tensor output_feats_tensor, int nProposal, int C) {
+  float *feats = feats_tensor.data_ptr<float>();
+  int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
+  float *output_feats = output_feats_tensor.data_ptr<float>();
+
+  global_avg_pool_fp_cuda(nProposal, C, feats, proposals_offset, output_feats);
+}
+
+void global_avg_pool_bp(at::Tensor d_feats_tensor,
+                        at::Tensor proposals_offset_tensor,
+                        at::Tensor d_output_feats_tensor, int nProposal,
+                        int C) {
+  float *d_feats = d_feats_tensor.data_ptr<float>();
+  int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
+  float *d_output_feats = d_output_feats_tensor.data_ptr<float>();
+
+  global_avg_pool_bp_cuda(nProposal, C, d_feats, proposals_offset,
+                          d_output_feats);
+}

softgroup/ops/src/roipool/roipool.cu

@@ -8,64 +8,6 @@ All Rights Reserved 2020.
 #include <math.h>
 #include <stdio.h>
 
-// fp
-__global__ void roipool_fp_cuda_(int nProposal, int C, float *feats,
-                                 int *proposals_offset, float *output_feats,
-                                 int *output_maxidx) {
-  for (int pp_id = blockIdx.x; pp_id < nProposal; pp_id += gridDim.x) {
-    int start = proposals_offset[pp_id];
-    int end = proposals_offset[pp_id + 1];
-
-    for (int plane = threadIdx.x; plane < C; plane += blockDim.x) {
-      int argmax_idx = -1;
-      float max_val = -1e50;
-
-      for (int i = start; i < end; i++) {
-        if (feats[i * C + plane] > max_val) {
-          argmax_idx = i;
-          max_val = feats[i * C + plane];
-        }
-      }
-      output_maxidx[pp_id * C + plane] = argmax_idx;
-      output_feats[pp_id * C + plane] = max_val;
-    }
-  }
-}
-
-// input: feats (sumNPoint, C) float
-// input: proposals_offset (nProposal + 1) int
-// output: output_feats (nProposal, C) float
-// output: output_maxidx (nProposal, C) int
-void roipool_fp_cuda(int nProposal, int C, float *feats, int *proposals_offset,
-                     float *output_feats, int *output_maxidx) {
-  roipool_fp_cuda_<<<std::min(nProposal, (int)32768), std::min(C, (int)32)>>>(
-      nProposal, C, feats, proposals_offset, output_feats, output_maxidx);
-}
-
-// bp
-__global__ void roipool_bp_cuda_(int nProposal, int C, float *d_feats,
-                                 int *proposals_offset, int *output_maxidx,
-                                 float *d_output_feats) {
-  for (int pp_id = blockIdx.x; pp_id < nProposal; pp_id += gridDim.x) {
-    for (int plane = threadIdx.x; plane < C; plane += blockDim.x) {
-      int argmax_idx = output_maxidx[pp_id * C + plane];
-      atomicAdd(&d_feats[argmax_idx * C + plane],
-                d_output_feats[pp_id * C + plane]);
-    }
-  }
-}
-
-// input: d_output_feats (nProposal, C) float
-// input: output_maxidx (nProposal, C) int
-// input: proposals_offset (nProposal + 1) int
-// output: d_feats (sumNPoint, C) float
-void roipool_bp_cuda(int nProposal, int C, float *d_feats,
-                     int *proposals_offset, int *output_maxidx,
-                     float *d_output_feats) {
-  roipool_bp_cuda_<<<std::min(nProposal, (int)32768), std::min(C, (int)32)>>>(
-      nProposal, C, d_feats, proposals_offset, output_maxidx, d_output_feats);
-}
-
 // fp
 __global__ void global_avg_pool_fp_cuda_(int nProposal, int C, float *feats,
                                          int *proposals_offset,

softgroup/ops/src/roipool/roipool.h

@@ -11,23 +11,6 @@ All Rights Reserved 2020.
 
 #include "../datatype/datatype.h"
 
-//
-void roipool_fp(at::Tensor feats_tensor, at::Tensor proposals_offset_tensor,
-                at::Tensor output_feats_tensor, at::Tensor output_maxidx_tensor,
-                int nProposal, int C);
-
-void roipool_fp_cuda(int nProposal, int C, float *feats, int *proposals_offset,
-                     float *output_feats, int *output_maxidx);
-
-//
-void roipool_bp(at::Tensor d_feats_tensor, at::Tensor proposals_offset_tensor,
-                at::Tensor output_maxidx_tensor,
-                at::Tensor d_output_feats_tensor, int nProposal, int C);
-
-void roipool_bp_cuda(int nProposal, int C, float *d_feats,
-                     int *proposals_offset, int *output_maxidx,
-                     float *d_output_feats);
-
 void global_avg_pool_fp_cuda(int nProposal, int C, float *feats,
                              int *proposals_offset, float *output_feats);
 

softgroup/ops/src/sec_mean/sec_mean.cpp

@@ -8,27 +8,27 @@ All Rights Reserved 2020.
 
 void sec_mean(at::Tensor inp_tensor, at::Tensor offsets_tensor,
               at::Tensor out_tensor, int nProposal, int C) {
-  int *offsets = offsets_tensor.data<int>();
+  int *offsets = offsets_tensor.data_ptr<int>();
-  float *inp = inp_tensor.data<float>();
+  float *inp = inp_tensor.data_ptr<float>();
-  float *out = out_tensor.data<float>();
+  float *out = out_tensor.data_ptr<float>();
 
   sec_mean_cuda(nProposal, C, inp, offsets, out);
 }
 
 void sec_min(at::Tensor inp_tensor, at::Tensor offsets_tensor,
              at::Tensor out_tensor, int nProposal, int C) {
-  int *offsets = offsets_tensor.data<int>();
+  int *offsets = offsets_tensor.data_ptr<int>();
-  float *inp = inp_tensor.data<float>();
+  float *inp = inp_tensor.data_ptr<float>();
-  float *out = out_tensor.data<float>();
+  float *out = out_tensor.data_ptr<float>();
 
   sec_min_cuda(nProposal, C, inp, offsets, out);
 }
 
 void sec_max(at::Tensor inp_tensor, at::Tensor offsets_tensor,
              at::Tensor out_tensor, int nProposal, int C) {
-  int *offsets = offsets_tensor.data<int>();
+  int *offsets = offsets_tensor.data_ptr<int>();
-  float *inp = inp_tensor.data<float>();
+  float *inp = inp_tensor.data_ptr<float>();
-  float *out = out_tensor.data<float>();
+  float *out = out_tensor.data_ptr<float>();
 
   sec_max_cuda(nProposal, C, inp, offsets, out);
 }

softgroup/ops/src/softgroup_api.cpp

@@ -13,20 +13,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("voxelize_idx", &voxelize_idx_3d, "voxelize_idx");
   m.def("voxelize_fp", &voxelize_fp_feat, "voxelize_fp");
   m.def("voxelize_bp", &voxelize_bp_feat, "voxelize_bp");
-  m.def("point_recover_fp", &point_recover_fp_feat, "point_recover_fp");
-  m.def("point_recover_bp", &point_recover_bp_feat, "point_recover_bp");
 
   m.def("ballquery_batch_p", &ballquery_batch_p, "ballquery_batch_p");
   m.def("bfs_cluster", &bfs_cluster, "bfs_cluster");
 
-  m.def("roipool_fp", &roipool_fp, "roipool_fp");
-  m.def("roipool_bp", &roipool_bp, "roipool_bp");
-
   m.def("global_avg_pool_fp", &global_avg_pool_fp, "global_avg_pool_fp");
   m.def("global_avg_pool_bp", &global_avg_pool_bp, "global_avg_pool_bp");
 
-  m.def("get_iou", &get_iou, "get_iou");
-
   m.def("sec_mean", &sec_mean, "sec_mean");
   m.def("sec_min", &sec_min, "sec_min");
   m.def("sec_max", &sec_max, "sec_max");

softgroup/ops/src/softgroup_ops.cpp

@@ -5,7 +5,6 @@
 #include "bfs_cluster/bfs_cluster.cpp"
 #include "cal_iou_and_masklabel/cal_iou_and_masklabel.cpp"
 #include "datatype/datatype.cpp"
-#include "get_iou/get_iou.cpp"
 #include "roipool/roipool.cpp"
 #include "sec_mean/sec_mean.cpp"
 #include "voxelize/voxelize.cpp"
@@ -35,19 +34,3 @@ void voxelize_bp_feat(/* cuda float M*C */ at::Tensor d_output_feats,
   voxelize_bp<float>(d_output_feats, d_feats, output_map, mode, nActive,
                      maxActive, nPlane);
 }
-
-void point_recover_fp_feat(/* cuda float M*C */ at::Tensor feats,
-                           /* cuda float N*C */ at::Tensor output_feats,
-                           /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                           Int nActive, Int maxActive, Int nPlane) {
-  point_recover_fp<float>(feats, output_feats, idx_map, nActive, maxActive,
-                          nPlane);
-}
-
-void point_recover_bp_feat(/* cuda float N*C */ at::Tensor d_output_feats,
-                           /* cuda float M*C */ at::Tensor d_feats,
-                           /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                           Int nActive, Int maxActive, Int nPlane) {
-  point_recover_bp<float>(d_output_feats, d_feats, idx_map, nActive, maxActive,
-                          nPlane);
-}

softgroup/ops/src/softgroup_ops.h

@@ -3,7 +3,6 @@
 #include "bfs_cluster/bfs_cluster.h"
 #include "cal_iou_and_masklabel/cal_iou_and_masklabel.h"
 #include "datatype/datatype.h"
-#include "get_iou/get_iou.h"
 #include "roipool/roipool.h"
 #include "sec_mean/sec_mean.h"
 
@@ -24,14 +23,4 @@ void voxelize_bp_feat(/* cuda float M*C */ at::Tensor d_output_feats,
                       /* cuda Int M*(maxActive+1) */ at::Tensor output_map,
                       Int mode, Int nActive, Int maxActive, Int nPlane);
 
-void point_recover_fp_feat(/* cuda float M*C */ at::Tensor feats,
-                           /* cuda float N*C */ at::Tensor output_feats,
-                           /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                           Int nActive, Int maxActive, Int nPlane);
-
-void point_recover_bp_feat(/* cuda float N*C */ at::Tensor d_output_feats,
-                           /* cuda float M*C */ at::Tensor d_feats,
-                           /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                           Int nActive, Int maxActive, Int nPlane);
-
 #endif // HAIS_H

softgroup/ops/src/voxelize/voxelize.cpp

@@ -23,8 +23,8 @@ void voxelize_idx(/* long N*4 */ at::Tensor coords,
   Int nActive = 0;
 
   Int maxActive = voxelize_inputmap<dimension>(
-      inputSGs, input_map.data<Int>(), voxelizeRuleBook, nActive,
+      inputSGs, input_map.data_ptr<Int>(), voxelizeRuleBook, nActive,
-      coords.data<long>(), coords.size(0), coords.size(1), batchSize, mode);
+      coords.data_ptr<long>(), coords.size(0), coords.size(1), batchSize, mode);
 
   output_map.resize_({nActive, maxActive + 1});
   output_map.zero_();
@@ -32,9 +32,9 @@ void voxelize_idx(/* long N*4 */ at::Tensor coords,
   output_coords.resize_({nActive, coords.size(1)});
   output_coords.zero_();
 
-  Int *oM = output_map.data<Int>();
+  Int *oM = output_map.data_ptr<Int>();
-  long *oC = output_coords.data<long>();
+  long *oC = output_coords.data_ptr<long>();
-  voxelize_outputmap<dimension>(coords.data<long>(), oC, oM,
+  voxelize_outputmap<dimension>(coords.data_ptr<long>(), oC, oM,
                                 &voxelizeRuleBook[1][0], nActive, maxActive);
 }
 
@@ -173,10 +173,10 @@ void voxelize_fp(
     /* cuda Int M*(maxActive+1) */ at::Tensor output_map, Int mode, Int nActive,
     Int maxActive, Int nPlane) {
 
-  auto iF = feats.data<T>();
+  auto iF = feats.data_ptr<T>();
-  auto oF = output_feats.data<T>();
+  auto oF = output_feats.data_ptr<T>();
 
-  Int *rules = output_map.data<Int>();
+  Int *rules = output_map.data_ptr<Int>();
 
   voxelize_fp_cuda<T>(nActive, maxActive, nPlane, iF, oF, rules, mode == 4);
 }
@@ -186,38 +186,10 @@ void voxelize_bp(/* cuda float M*C */ at::Tensor d_output_feats,
                  /* cuda float N*C */ at::Tensor d_feats,
                  /* cuda Int M*(maxActive+1) */ at::Tensor output_map, Int mode,
                  Int nActive, Int maxActive, Int nPlane) {
-  auto d_oF = d_output_feats.data<T>();
+  auto d_oF = d_output_feats.data_ptr<T>();
-  auto d_iF = d_feats.data<T>();
+  auto d_iF = d_feats.data_ptr<T>();
 
-  Int *rules = output_map.data<Int>();
+  Int *rules = output_map.data_ptr<Int>();
 
   voxelize_bp_cuda<T>(nActive, maxActive, nPlane, d_oF, d_iF, rules, mode == 4);
 }
-
-/* ================================== point_recover
- * ================================== */
-template <typename T>
-void point_recover_fp(/* cuda float M*C */ at::Tensor feats,
-                      /* cuda float N*C */ at::Tensor output_feats,
-                      /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                      Int nActive, Int maxActive, Int nPlane) {
-  auto iF = feats.data<T>();
-  auto oF = output_feats.data<T>();
-
-  Int *rules = idx_map.data<Int>();
-
-  voxelize_bp_cuda<T>(nActive, maxActive, nPlane, iF, oF, rules, false);
-}
-
-template <typename T>
-void point_recover_bp(/* cuda float N*C */ at::Tensor d_output_feats,
-                      /* cuda float M*C */ at::Tensor d_feats,
-                      /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                      Int nActive, Int maxActive, Int nPlane) {
-  auto d_oF = d_output_feats.data<T>();
-  auto d_iF = d_feats.data<T>();
-
-  Int *rules = idx_map.data<Int>();
-
-  voxelize_fp_cuda<T>(nActive, maxActive, nPlane, d_oF, d_iF, rules, false);
-}

softgroup/ops/src/voxelize/voxelize.h

@@ -54,19 +54,4 @@ template <typename T>
 void voxelize_bp_cuda(Int nOutputRows, Int maxActive, Int nPlanes,
                       T *d_output_feats, T *d_feats, Int *rules, bool average);
 
-/* ================================== point_recover
- * ================================== */
-template <typename T>
-void point_recover_fp(/* cuda float M*C */ at::Tensor feats,
-                      /* cuda float N*C */ at::Tensor output_feats,
-                      /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                      Int nActive, Int maxActive, Int nPlane);
-
-//
-template <typename T>
-void point_recover_bp(/* cuda float N*C */ at::Tensor d_output_feats,
-                      /* cuda float M*C */ at::Tensor d_feats,
-                      /* cuda Int M*(maxActive+1) */ at::Tensor idx_map,
-                      Int nActive, Int maxActive, Int nPlane);
-
 #endif // VOXELIZE_H

tools/convert_checkpoint.py

@@ -0,0 +1,29 @@
+# Convert spconv1 checkpoint to spconv2 checkpoint
+import argparse
+from collections import OrderedDict
+
+import torch
+
+parser = argparse.ArgumentParser()
+parser.add_argument('checkpoint', type=str, help='spconv1 checkpoint')
+args = parser.parse_args()
+
+checkpoint = torch.load(args.checkpoint)
+model = checkpoint['net']
+new_model = OrderedDict()
+
+for k, v in model.items():
+    new_k, new_v = k, v
+    if 'weight' in k and len(v.size()) == 5:
+        # KKKIO to OKKKI (0, 1, 2, 3, 4) -> (4, 0, 1, 2, 3)
+        new_v = v.permute(4, 0, 1, 2, 3)
+    if 'intra_ins_unet' in k:
+        new_k = k.replace('intra_ins_unet', 'tiny_unet')
+    elif 'score_linear' in new_k:
+        new_k = k.replace('score_linear', 'iou_score_linear')
+    elif 'intra_ins_outputlayer' in k:
+        new_k = k.replace('intra_ins_outputlayer', 'tiny_unet_outputlayer')
+    new_model[new_k] = new_v
+
+checkpoint['net'] = new_model
+torch.save(checkpoint, args.checkpoint.replace('.pth', '_spconv2.pth'))

tools/dist_test.sh

@@ -4,4 +4,4 @@ CHECK_POINT=$2
 GPUS=$3
 PORT=${PORT:-29501}
 
-OMP_NUM_THREADS=1 torchrun --nproc_per_node=$GPUS --master_port=$PORT ./test.py $CONFIG $CHECK_POINT --dist ${@:4}
+OMP_NUM_THREADS=1 torchrun --nproc_per_node=$GPUS --master_port=$PORT  $(dirname "$0")/test.py $CONFIG $CHECK_POINT --dist ${@:4}

tools/dist_train.sh

@@ -3,4 +3,4 @@ CONFIG=$1
 GPUS=$2
 PORT=${PORT:-29500}
 
-OMP_NUM_THREADS=1 torchrun --nproc_per_node=$GPUS --master_port=$PORT ./train.py --dist $CONFIG ${@:3}
+OMP_NUM_THREADS=1 torchrun --nproc_per_node=$GPUS --master_port=$PORT $(dirname "$0")/train.py --dist $CONFIG ${@:3}