vhad/SoftGroup
1
0
Fork 0
mirror of https://github.com/botastic/SoftGroup.git synced 2025-10-16 11:45:42 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add setup.py

Browse Source
This commit is contained in:
Thang Vu 2022-04-15 08:16:12 +00:00
parent 3b39cc3bf1
commit 590b96b8aa
43 changed files with 160 additions and 498 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
requirements.txt
View File

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

26
setup.py Normal file
View File

@ -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})

2
softgroup/data/custom.py
View File

@ -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):

2
softgroup/data/s3dis.py
View File

@ -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

1
softgroup/lib/softgroup_ops/__init__.py
View File

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

23
softgroup/lib/softgroup_ops/src/get_iou/get_iou.cpp
View File

@ -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);
}

48
softgroup/lib/softgroup_ops/src/get_iou/get_iou.cu
View File

@ -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);
}

24
softgroup/lib/softgroup_ops/src/get_iou/get_iou.h
View File

@ -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

53
softgroup/lib/softgroup_ops/src/roipool/roipool.cpp
View File

@ -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);
}

6
softgroup/model/softgroup.py
View File

@ -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

1
softgroup/ops/__init__.py Normal file
View File

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

0
softgroup/lib/softgroup_ops/clang_format.sh → softgroup/ops/clang_format.sh
View File

159
softgroup/lib/softgroup_ops/functions/softgroup_ops.py → softgroup/ops/functions.py
View File

@ -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

0
softgroup/lib/softgroup_ops/setup.py → softgroup/ops/setup.py
View File

20
softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cpp → softgroup/ops/src/bfs_cluster/bfs_cluster.cpp
View File

@ -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 int *batch_idxs = batch_idxs_tensor.data<int>();
const int *batch_offsets = batch_offsets_tensor.data<int>();
int *idx = idx_tensor.data<int>();
int *start_len = start_len_tensor.data<int>();
const float *xyz = xyz_tensor.data_ptr<float>();
const int *batch_idxs = batch_idxs_tensor.data_ptr<int>();
const int *batch_offsets = batch_offsets_tensor.data_ptr<int>();
int *idx = idx_tensor.data_ptr<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>();
Int *ball_query_idxs = ball_query_idxs_tensor.data<Int>();
int *start_len = start_len_tensor.data<int>();
float *class_numpoint_mean = class_numpoint_mean_tensor.data_ptr<float>();
Int *ball_query_idxs = ball_query_idxs_tensor.data_ptr<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_offsets = cluster_offsets_tensor.data<int>();
int *cluster_idxs = cluster_idxs_tensor.data_ptr<int>();
int *cluster_offsets = cluster_offsets_tensor.data_ptr<int>();
fill_cluster_idxs_(CCs, cluster_idxs, cluster_offsets);
}

0
softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.cu → softgroup/ops/src/bfs_cluster/bfs_cluster.cu
View File

1
softgroup/lib/softgroup_ops/src/bfs_cluster/bfs_cluster.h → softgroup/ops/src/bfs_cluster/bfs_cluster.h
View File

@ -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"

34
softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp → softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cpp
View File

@ -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_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>();
int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
long *instance_labels = instance_labels_tensor.data_ptr<long>();
int *instance_pointnum = instance_pointnum_tensor.data_ptr<int>();
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_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>();
float *mask_scores_sigmoid = mask_scores_sigmoid_tensor.data<float>();
int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
long *instance_labels = instance_labels_tensor.data_ptr<long>();
int *instance_pointnum = instance_pointnum_tensor.data_ptr<int>();
float *proposals_iou = proposals_iou_tensor.data_ptr<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_offset = proposals_offset_tensor.data<int>();
long *instance_labels = instance_labels_tensor.data<long>();
long *instance_cls = instance_cls_tensor.data<long>();
float *proposals_iou = proposals_iou_tensor.data<float>();
float *mask_label = mask_labels_tensor.data<float>();
int *proposals_idx = proposals_idx_tensor.data_ptr<int>();
int *proposals_offset = proposals_offset_tensor.data_ptr<int>();
long *instance_labels = instance_labels_tensor.data_ptr<long>();
long *instance_cls = instance_cls_tensor.data_ptr<long>();
float *proposals_iou = proposals_iou_tensor.data_ptr<float>();
float *mask_label = mask_labels_tensor.data_ptr<float>();
// input: nInstance (1,), int
// input: nProposal (1,), int

0
softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu → softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.cu
View File

0
softgroup/lib/softgroup_ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h → softgroup/ops/src/cal_iou_and_masklabel/cal_iou_and_masklabel.h
View File

1
softgroup/lib/softgroup_ops/src/cuda.cu → softgroup/ops/src/cuda.cu
View File

@ -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"

0
softgroup/lib/softgroup_ops/src/cuda_utils.h → softgroup/ops/src/cuda_utils.h
View File

0
softgroup/lib/softgroup_ops/src/datatype/datatype.cpp → softgroup/ops/src/datatype/datatype.cpp
View File

0
softgroup/lib/softgroup_ops/src/datatype/datatype.h → softgroup/ops/src/datatype/datatype.h
View File

29
softgroup/ops/src/roipool/roipool.cpp Normal file
View File

@ -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);
}

58
softgroup/lib/softgroup_ops/src/roipool/roipool.cu → softgroup/ops/src/roipool/roipool.cu
View File

@ -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,

17
softgroup/lib/softgroup_ops/src/roipool/roipool.h → softgroup/ops/src/roipool/roipool.h
View File

@ -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);

18
softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cpp → softgroup/ops/src/sec_mean/sec_mean.cpp
View File

@ -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>();
float *inp = inp_tensor.data<float>();
float *out = out_tensor.data<float>();
int *offsets = offsets_tensor.data_ptr<int>();
float *inp = inp_tensor.data_ptr<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>();
float *inp = inp_tensor.data<float>();
float *out = out_tensor.data<float>();
int *offsets = offsets_tensor.data_ptr<int>();
float *inp = inp_tensor.data_ptr<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>();
float *inp = inp_tensor.data<float>();
float *out = out_tensor.data<float>();
int *offsets = offsets_tensor.data_ptr<int>();
float *inp = inp_tensor.data_ptr<float>();
float *out = out_tensor.data_ptr<float>();
sec_max_cuda(nProposal, C, inp, offsets, out);
}

0
softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.cu → softgroup/ops/src/sec_mean/sec_mean.cu
View File

0
softgroup/lib/softgroup_ops/src/sec_mean/sec_mean.h → softgroup/ops/src/sec_mean/sec_mean.h
View File

7
softgroup/lib/softgroup_ops/src/softgroup_api.cpp → softgroup/ops/src/softgroup_api.cpp
View File

@ -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");

17
softgroup/lib/softgroup_ops/src/softgroup_ops.cpp → softgroup/ops/src/softgroup_ops.cpp
View File

@ -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);
}

11
softgroup/lib/softgroup_ops/src/softgroup_ops.h → softgroup/ops/src/softgroup_ops.h
View File

@ -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

50
softgroup/lib/softgroup_ops/src/voxelize/voxelize.cpp → softgroup/ops/src/voxelize/voxelize.cpp
View File

@ -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,
coords.data<long>(), coords.size(0), coords.size(1), batchSize, mode);
inputSGs, input_map.data_ptr<Int>(), voxelizeRuleBook, nActive,
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>();
long *oC = output_coords.data<long>();
voxelize_outputmap<dimension>(coords.data<long>(), oC, oM,
Int *oM = output_map.data_ptr<Int>();
long *oC = output_coords.data_ptr<long>();
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 oF = output_feats.data<T>();
auto iF = feats.data_ptr<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_iF = d_feats.data<T>();
auto d_oF = d_output_feats.data_ptr<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);
}

0
softgroup/lib/softgroup_ops/src/voxelize/voxelize.cu → softgroup/ops/src/voxelize/voxelize.cu
View File

15
softgroup/lib/softgroup_ops/src/voxelize/voxelize.h → softgroup/ops/src/voxelize/voxelize.h
View File

@ -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

29
tools/convert_checkpoint.py Normal file
View File

@ -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'))

2
dist_test.sh → tools/dist_test.sh
View File

@ -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}

2
dist_train.sh → tools/dist_train.sh
View File

@ -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}

0
eval_det.py → tools/eval_det.py
View File

0
test.py → tools/test.py
View File

0
train.py → tools/train.py
View File

0
visualization.py → tools/visualization.py
View File