diff --git a/requirements.txt b/requirements.txt index c184ea4..154fd3e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ munch pandas plyfile -pyyaml==5.4.1 +pyyaml scikit-learn scipy six diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..e42c0aa --- /dev/null +++ b/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}) diff --git a/softgroup/data/custom.py b/softgroup/data/custom.py index 3818ebd..fcda1bd 100644 --- a/softgroup/data/custom.py +++ b/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): diff --git a/softgroup/data/s3dis.py b/softgroup/data/s3dis.py index 63fd9f5..718b063 100644 --- a/softgroup/data/s3dis.py +++ b/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 diff --git a/softgroup/lib/softgroup_ops/__init__.py b/softgroup/lib/softgroup_ops/__init__.py deleted file mode 100644 index 74eca6b..0000000 --- a/softgroup/lib/softgroup_ops/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .functions.softgroup_ops import * diff --git a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cpp b/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cpp deleted file mode 100644 index 0795db5..0000000 --- a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cpp +++ /dev/null @@ -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 *proposals_offset = proposals_offset_tensor.data(); - long *instance_labels = instance_labels_tensor.data(); - int *instance_pointnum = instance_pointnum_tensor.data(); - - float *proposals_iou = proposals_iou_tensor.data(); - - get_iou_cuda(nInstance, nProposal, proposals_idx, proposals_offset, - instance_labels, instance_pointnum, proposals_iou); -} diff --git a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cu b/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cu deleted file mode 100644 index 96c3c86..0000000 --- a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.cu +++ /dev/null @@ -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 -#include - -__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_<<>>( - nInstance, nProposal, proposals_idx, proposals_offset, instance_labels, - instance_pointnum, proposals_iou); -} diff --git a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.h b/softgroup/lib/softgroup_ops/src/get_iou/get_iou.h deleted file mode 100644 index 5300065..0000000 --- a/softgroup/lib/softgroup_ops/src/get_iou/get_iou.h +++ /dev/null @@ -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 -#include - -#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 diff --git a/softgroup/lib/softgroup_ops/src/roipool/roipool.cpp b/softgroup/lib/softgroup_ops/src/roipool/roipool.cpp deleted file mode 100644 index f51c700..0000000 --- a/softgroup/lib/softgroup_ops/src/roipool/roipool.cpp +++ /dev/null @@ -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(); - int *proposals_offset = proposals_offset_tensor.data(); - float *output_feats = output_feats_tensor.data(); - int *output_maxidx = output_maxidx_tensor.data(); - - 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(); - int *proposals_offset = proposals_offset_tensor.data(); - int *output_maxidx = output_maxidx_tensor.data(); - float *d_output_feats = d_output_feats_tensor.data(); - - 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(); - int *proposals_offset = proposals_offset_tensor.data(); - float *output_feats = output_feats_tensor.data(); - - 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(); - int *proposals_offset = proposals_offset_tensor.data(); - float *d_output_feats = d_output_feats_tensor.data(); - - global_avg_pool_bp_cuda(nProposal, C, d_feats, proposals_offset, - d_output_feats); -} diff --git a/softgroup/model/softgroup.py b/softgroup/model/softgroup.py index 54007f6..ce44ae6 100644 --- a/softgroup/model/softgroup.py +++ b/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, - get_mask_iou_on_pred, get_mask_label, global_avg_pool, sec_max, - sec_min, voxelization, voxelization_idx) +from ..ops import (ballquery_batch_p, bfs_cluster, get_mask_iou_on_cluster, get_mask_iou_on_pred, + get_mask_label, global_avg_pool, sec_max, sec_min, voxelization, + voxelization_idx) from ..util import cuda_cast, force_fp32, rle_encode from .blocks import MLP, ResidualBlock, UBlock diff --git a/softgroup/ops/__init__.py b/softgroup/ops/__init__.py new file mode 100644 index 0000000..c8f7548 --- /dev/null +++ b/softgroup/ops/__init__.py @@ -0,0 +1 @@ +from .functions import * diff --git a/softgroup/lib/softgroup_ops/clang_format.sh b/softgroup/ops/clang_format.sh similarity index 100% rename from softgroup/lib/softgroup_ops/clang_format.sh rename to softgroup/ops/clang_format.sh diff --git a/softgroup/lib/softgroup_ops/functions/softgroup_ops.py b/softgroup/ops/functions.py similarity index 65% rename from softgroup/lib/softgroup_ops/functions/softgroup_ops.py rename to softgroup/ops/functions.py index 81eae86..bd18d94 100644 --- a/softgroup/lib/softgroup_ops/functions/softgroup_ops.py +++ b/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,8 +30,8 @@ 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, - instance_pointnum, proposals_iou, nInstance, nProposal) + ops.get_mask_iou_on_cluster(proposals_idx, proposals_offset, instance_labels, + instance_pointnum, proposals_iou, nInstance, nProposal) return proposals_iou @@ -71,9 +71,9 @@ 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, - instance_pointnum, proposals_iou, nInstance, nProposal, - mask_scores_sigmoid) + ops.get_mask_iou_on_pred(proposals_idx, proposals_offset, instance_labels, + instance_pointnum, proposals_iou, nInstance, nProposal, + mask_scores_sigmoid) return proposals_iou @@ -112,8 +112,8 @@ 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, - proposals_iou, nInstance, nProposal, iou_thr, mask_label) + 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, - C) + ops.voxelize_bp(d_output_feats.contiguous(), d_feats, map_rule, mode, M, maxActive, 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, - start_len, n, meanActive, radius) + nActive = ops.ballquery_batch_p(coords, batch_idxs, batch_offsets, idx, start_len, n, + meanActive, radius) if nActive <= n * meanActive: break meanActive = int(nActive // n + 1) @@ -296,8 +255,8 @@ 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, - cluster_offsets, N, threshold, class_id) + 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(), - nProposal, C) + ops.global_avg_pool_bp(d_feats, proposals_offset, d_output_feats.contiguous(), 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 diff --git a/softgroup/lib/softgroup_ops/setup.py b/softgroup/ops/setup.py similarity index 100% rename from softgroup/lib/softgroup_ops/setup.py rename to softgroup/ops/setup.py diff --git a/softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cpp b/softgroup/ops/src/bfs_cluster/bfs_cluster.cpp similarity index 87% rename from softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cpp rename to softgroup/ops/src/bfs_cluster/bfs_cluster.cpp index 26a1502..35a1c14 100644 --- a/softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cpp +++ b/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(); - const int *batch_idxs = batch_idxs_tensor.data(); - const int *batch_offsets = batch_offsets_tensor.data(); - int *idx = idx_tensor.data(); - int *start_len = start_len_tensor.data(); + const float *xyz = xyz_tensor.data_ptr(); + const int *batch_idxs = batch_idxs_tensor.data_ptr(); + const int *batch_offsets = batch_offsets_tensor.data_ptr(); + int *idx = idx_tensor.data_ptr(); + int *start_len = start_len_tensor.data_ptr(); 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(); - Int *ball_query_idxs = ball_query_idxs_tensor.data(); - int *start_len = start_len_tensor.data(); + float *class_numpoint_mean = class_numpoint_mean_tensor.data_ptr(); + Int *ball_query_idxs = ball_query_idxs_tensor.data_ptr(); + int *start_len = start_len_tensor.data_ptr(); 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 *cluster_offsets = cluster_offsets_tensor.data(); + int *cluster_idxs = cluster_idxs_tensor.data_ptr(); + int *cluster_offsets = cluster_offsets_tensor.data_ptr(); fill_cluster_idxs_(CCs, cluster_idxs, cluster_offsets); } diff --git a/softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cu b/softgroup/ops/src/bfs_cluster/bfs_cluster.cu similarity index 100% rename from softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cu rename to softgroup/ops/src/bfs_cluster/bfs_cluster.cu diff --git a/softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.h b/softgroup/ops/src/bfs_cluster/bfs_cluster.h similarity index 98% rename from softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.h rename to softgroup/ops/src/bfs_cluster/bfs_cluster.h index 37ca6b7..1e088b5 100644 --- a/softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.h +++ b/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 -#include #include #include "../datatype/datatype.h" diff --git a/softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp b/softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp similarity index 73% rename from softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp rename to softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp index 06916fa..b1b7531 100644 --- a/softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp +++ b/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 *proposals_offset = proposals_offset_tensor.data(); - long *instance_labels = instance_labels_tensor.data(); - int *instance_pointnum = instance_pointnum_tensor.data(); - float *proposals_iou = proposals_iou_tensor.data(); + int *proposals_idx = proposals_idx_tensor.data_ptr(); + int *proposals_offset = proposals_offset_tensor.data_ptr(); + long *instance_labels = instance_labels_tensor.data_ptr(); + int *instance_pointnum = instance_pointnum_tensor.data_ptr(); + float *proposals_iou = proposals_iou_tensor.data_ptr(); // 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 *proposals_offset = proposals_offset_tensor.data(); - long *instance_labels = instance_labels_tensor.data(); - int *instance_pointnum = instance_pointnum_tensor.data(); - float *proposals_iou = proposals_iou_tensor.data(); - float *mask_scores_sigmoid = mask_scores_sigmoid_tensor.data(); + int *proposals_idx = proposals_idx_tensor.data_ptr(); + int *proposals_offset = proposals_offset_tensor.data_ptr(); + long *instance_labels = instance_labels_tensor.data_ptr(); + int *instance_pointnum = instance_pointnum_tensor.data_ptr(); + float *proposals_iou = proposals_iou_tensor.data_ptr(); + float *mask_scores_sigmoid = mask_scores_sigmoid_tensor.data_ptr(); // 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 *proposals_offset = proposals_offset_tensor.data(); - long *instance_labels = instance_labels_tensor.data(); - long *instance_cls = instance_cls_tensor.data(); - float *proposals_iou = proposals_iou_tensor.data(); - float *mask_label = mask_labels_tensor.data(); + int *proposals_idx = proposals_idx_tensor.data_ptr(); + int *proposals_offset = proposals_offset_tensor.data_ptr(); + long *instance_labels = instance_labels_tensor.data_ptr(); + long *instance_cls = instance_cls_tensor.data_ptr(); + float *proposals_iou = proposals_iou_tensor.data_ptr(); + float *mask_label = mask_labels_tensor.data_ptr(); // input: nInstance (1,), int // input: nProposal (1,), int diff --git a/softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu b/softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu similarity index 100% rename from softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu rename to softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu diff --git a/softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h b/softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h similarity index 100% rename from softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h rename to softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h diff --git a/softgroup/lib/softgroup_ops/src/cuda.cu b/softgroup/ops/src/cuda.cu similarity index 96% rename from softgroup/lib/softgroup_ops/src/cuda.cu rename to softgroup/ops/src/cuda.cu index b611e88..9767e58 100644 --- a/softgroup/lib/softgroup_ops/src/cuda.cu +++ b/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" diff --git a/softgroup/lib/softgroup_ops/src/cuda_utils.h b/softgroup/ops/src/cuda_utils.h similarity index 100% rename from softgroup/lib/softgroup_ops/src/cuda_utils.h rename to softgroup/ops/src/cuda_utils.h diff --git a/softgroup/lib/softgroup_ops/src/datatype/datatype.cpp b/softgroup/ops/src/datatype/datatype.cpp similarity index 100% rename from softgroup/lib/softgroup_ops/src/datatype/datatype.cpp rename to softgroup/ops/src/datatype/datatype.cpp diff --git a/softgroup/lib/softgroup_ops/src/datatype/datatype.h b/softgroup/ops/src/datatype/datatype.h similarity index 100% rename from softgroup/lib/softgroup_ops/src/datatype/datatype.h rename to softgroup/ops/src/datatype/datatype.h diff --git a/softgroup/ops/src/roipool/roipool.cpp b/softgroup/ops/src/roipool/roipool.cpp new file mode 100644 index 0000000..7e9d998 --- /dev/null +++ b/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(); + int *proposals_offset = proposals_offset_tensor.data_ptr(); + float *output_feats = output_feats_tensor.data_ptr(); + + 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(); + int *proposals_offset = proposals_offset_tensor.data_ptr(); + float *d_output_feats = d_output_feats_tensor.data_ptr(); + + global_avg_pool_bp_cuda(nProposal, C, d_feats, proposals_offset, + d_output_feats); +} diff --git a/softgroup/lib/softgroup_ops/src/roipool/roipool.cu b/softgroup/ops/src/roipool/roipool.cu similarity index 52% rename from softgroup/lib/softgroup_ops/src/roipool/roipool.cu rename to softgroup/ops/src/roipool/roipool.cu index a9d2021..10c60aa 100644 --- a/softgroup/lib/softgroup_ops/src/roipool/roipool.cu +++ b/softgroup/ops/src/roipool/roipool.cu @@ -8,64 +8,6 @@ All Rights Reserved 2020. #include #include -// 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_<<>>( - 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_<<>>( - 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, diff --git a/softgroup/lib/softgroup_ops/src/roipool/roipool.h b/softgroup/ops/src/roipool/roipool.h similarity index 55% rename from softgroup/lib/softgroup_ops/src/roipool/roipool.h rename to softgroup/ops/src/roipool/roipool.h index 9877ced..bfe7ddc 100644 --- a/softgroup/lib/softgroup_ops/src/roipool/roipool.h +++ b/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); diff --git a/softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cpp b/softgroup/ops/src/sec_mean/sec_mean.cpp similarity index 60% rename from softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cpp rename to softgroup/ops/src/sec_mean/sec_mean.cpp index ff659db..3c08479 100644 --- a/softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cpp +++ b/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(); - float *inp = inp_tensor.data(); - float *out = out_tensor.data(); + int *offsets = offsets_tensor.data_ptr(); + float *inp = inp_tensor.data_ptr(); + float *out = out_tensor.data_ptr(); 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(); - float *inp = inp_tensor.data(); - float *out = out_tensor.data(); + int *offsets = offsets_tensor.data_ptr(); + float *inp = inp_tensor.data_ptr(); + float *out = out_tensor.data_ptr(); 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(); - float *inp = inp_tensor.data(); - float *out = out_tensor.data(); + int *offsets = offsets_tensor.data_ptr(); + float *inp = inp_tensor.data_ptr(); + float *out = out_tensor.data_ptr(); sec_max_cuda(nProposal, C, inp, offsets, out); } diff --git a/softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cu b/softgroup/ops/src/sec_mean/sec_mean.cu similarity index 100% rename from softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cu rename to softgroup/ops/src/sec_mean/sec_mean.cu diff --git a/softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.h b/softgroup/ops/src/sec_mean/sec_mean.h similarity index 100% rename from softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.h rename to softgroup/ops/src/sec_mean/sec_mean.h diff --git a/softgroup/lib/softgroup_ops/src/softgroup_api.cpp b/softgroup/ops/src/softgroup_api.cpp similarity index 76% rename from softgroup/lib/softgroup_ops/src/softgroup_api.cpp rename to softgroup/ops/src/softgroup_api.cpp index 5e05f7d..4044e0b 100644 --- a/softgroup/lib/softgroup_ops/src/softgroup_api.cpp +++ b/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"); diff --git a/softgroup/lib/softgroup_ops/src/softgroup_ops.cpp b/softgroup/ops/src/softgroup_ops.cpp similarity index 63% rename from softgroup/lib/softgroup_ops/src/softgroup_ops.cpp rename to softgroup/ops/src/softgroup_ops.cpp index dc71a19..9979a96 100644 --- a/softgroup/lib/softgroup_ops/src/softgroup_ops.cpp +++ b/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(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(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(d_output_feats, d_feats, idx_map, nActive, maxActive, - nPlane); -} diff --git a/softgroup/lib/softgroup_ops/src/softgroup_ops.h b/softgroup/ops/src/softgroup_ops.h similarity index 63% rename from softgroup/lib/softgroup_ops/src/softgroup_ops.h rename to softgroup/ops/src/softgroup_ops.h index 0cd0349..8a0b1fa 100644 --- a/softgroup/lib/softgroup_ops/src/softgroup_ops.h +++ b/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 diff --git a/softgroup/lib/softgroup_ops/src/voxelize/voxelize.cpp b/softgroup/ops/src/voxelize/voxelize.cpp similarity index 78% rename from softgroup/lib/softgroup_ops/src/voxelize/voxelize.cpp rename to softgroup/ops/src/voxelize/voxelize.cpp index 8db7c2c..6664072 100644 --- a/softgroup/lib/softgroup_ops/src/voxelize/voxelize.cpp +++ b/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( - inputSGs, input_map.data(), voxelizeRuleBook, nActive, - coords.data(), coords.size(0), coords.size(1), batchSize, mode); + inputSGs, input_map.data_ptr(), voxelizeRuleBook, nActive, + coords.data_ptr(), 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(); - long *oC = output_coords.data(); - voxelize_outputmap(coords.data(), oC, oM, + Int *oM = output_map.data_ptr(); + long *oC = output_coords.data_ptr(); + voxelize_outputmap(coords.data_ptr(), 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(); - auto oF = output_feats.data(); + auto iF = feats.data_ptr(); + auto oF = output_feats.data_ptr(); - Int *rules = output_map.data(); + Int *rules = output_map.data_ptr(); voxelize_fp_cuda(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(); - auto d_iF = d_feats.data(); + auto d_oF = d_output_feats.data_ptr(); + auto d_iF = d_feats.data_ptr(); - Int *rules = output_map.data(); + Int *rules = output_map.data_ptr(); voxelize_bp_cuda(nActive, maxActive, nPlane, d_oF, d_iF, rules, mode == 4); } - -/* ================================== point_recover - * ================================== */ -template -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(); - auto oF = output_feats.data(); - - Int *rules = idx_map.data(); - - voxelize_bp_cuda(nActive, maxActive, nPlane, iF, oF, rules, false); -} - -template -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(); - auto d_iF = d_feats.data(); - - Int *rules = idx_map.data(); - - voxelize_fp_cuda(nActive, maxActive, nPlane, d_oF, d_iF, rules, false); -} diff --git a/softgroup/lib/softgroup_ops/src/voxelize/voxelize.cu b/softgroup/ops/src/voxelize/voxelize.cu similarity index 100% rename from softgroup/lib/softgroup_ops/src/voxelize/voxelize.cu rename to softgroup/ops/src/voxelize/voxelize.cu diff --git a/softgroup/lib/softgroup_ops/src/voxelize/voxelize.h b/softgroup/ops/src/voxelize/voxelize.h similarity index 75% rename from softgroup/lib/softgroup_ops/src/voxelize/voxelize.h rename to softgroup/ops/src/voxelize/voxelize.h index 04b43fb..df7ce2d 100644 --- a/softgroup/lib/softgroup_ops/src/voxelize/voxelize.h +++ b/softgroup/ops/src/voxelize/voxelize.h @@ -54,19 +54,4 @@ template void voxelize_bp_cuda(Int nOutputRows, Int maxActive, Int nPlanes, T *d_output_feats, T *d_feats, Int *rules, bool average); -/* ================================== point_recover - * ================================== */ -template -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 -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 diff --git a/tools/convert_checkpoint.py b/tools/convert_checkpoint.py new file mode 100644 index 0000000..80d2e79 --- /dev/null +++ b/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')) diff --git a/dist_test.sh b/tools/dist_test.sh similarity index 67% rename from dist_test.sh rename to tools/dist_test.sh index 3130489..fa1e457 100755 --- a/dist_test.sh +++ b/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} diff --git a/dist_train.sh b/tools/dist_train.sh similarity index 69% rename from dist_train.sh rename to tools/dist_train.sh index d870d1f..08d6df0 100755 --- a/dist_train.sh +++ b/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} diff --git a/eval_det.py b/tools/eval_det.py similarity index 100% rename from eval_det.py rename to tools/eval_det.py diff --git a/test.py b/tools/test.py similarity index 100% rename from test.py rename to tools/test.py diff --git a/train.py b/tools/train.py similarity index 100% rename from train.py rename to tools/train.py diff --git a/visualization.py b/tools/visualization.py similarity index 100% rename from visualization.py rename to tools/visualization.py