wassname/Pointnet2_PyTorch
1
0
Fork 0
mirror of https://github.com/wassname/Pointnet2_PyTorch.git synced 2026-06-27 16:00:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Updates

Browse Source
This commit is contained in:
erikwijmans
2018-01-06 12:13:52 -05:00
parent 7e746ba72a
commit 5a5adc2b77
20 changed files with 650 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+2 -2
View File
@@ -11,7 +11,7 @@ cuda_compile(cuda_kernels SHARED ${cuda_kernels_src} OPTIONS -O3)
set(BUILD_CMD python "${CMAKE_CURRENT_SOURCE_DIR}/utils/build_ffi.py")
file(GLOB wrapper_headers "${CMAKE_CURRENT_SOURCE_DIR}/utils/cinclude/*wrapper.h")
file(GLOB wrapper_sources "${CMAKE_CURRENT_SOURCE_DIR}/utils/csrs/*.c")
add_custom_command(OUTPUT "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext/__ext.so"
add_custom_command(OUTPUT "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext/pointnet2/_pointnet2.so"
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/utils
COMMAND ${BUILD_CMD} --build --objs ${cuda_kernels}
DEPENDS ${cuda_kernels}
@@ -20,6 +20,6 @@ add_custom_command(OUTPUT "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext/__ext.so"
VERBATIM)
add_custom_target(ext ALL
DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext/__ext.so")
DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext/pointnet2/_pointnet2.so")
set_directory_properties(PROPERTIES ADDITIONAL_MAKE_CLEAN_FILES "${CMAKE_CURRENT_SOURCE_DIR}/utils/_ext")
data/Indoor3DSemSegLoader.py
+17 -12
View File
@@ -16,12 +16,10 @@ def _load_data_file(name):
class Indoor3DSemSeg(data.Dataset):
def __init__(self,
num_points,
root,
train=True,
download=True,
data_precent=1.0):
def __init__(
self, num_points, root, train=True, download=True, data_precent=1.0
):
super().__init__()
self.data_precent = data_precent
root = os.path.abspath(root)
@@ -32,18 +30,23 @@ class Indoor3DSemSeg(data.Dataset):
if download and not os.path.exists(self.data_dir):
zipfile = os.path.join(root, os.path.basename(self.url))
subprocess.check_call(
shlex.split("curl {} -o {}".format(self.url, zipfile)))
shlex.split("curl {} -o {}".format(self.url, zipfile))
)
subprocess.check_call(shlex.split("unzip {} -d {}".format(zipfile, root)))
subprocess.check_call(
shlex.split("unzip {} -d {}".format(zipfile, root))
)
subprocess.check_call(shlex.split("rm {}".format(zipfile)))
self.train, self.num_points = train, num_points
all_files = _get_data_files(
os.path.join(self.data_dir, "all_files.txt"))
os.path.join(self.data_dir, "all_files.txt")
)
room_filelist = _get_data_files(
os.path.join(self.data_dir, "room_filelist.txt"))
os.path.join(self.data_dir, "room_filelist.txt")
)
data_batchlist, label_batchlist = [], []
for f in all_files:
@@ -74,9 +77,11 @@ class Indoor3DSemSeg(data.Dataset):
np.random.shuffle(pt_idxs)
current_points = torch.from_numpy(self.points[idx, pt_idxs, :]).type(
torch.FloatTensor)
torch.FloatTensor
)
current_labels = torch.from_numpy(self.labels[idx, pt_idxs]).type(
torch.LongTensor)
torch.LongTensor
)
return current_points, current_labels
data/ModelNet40Loader.py
+13 -11
View File
@@ -19,12 +19,10 @@ def _load_data_file(name):
class ModelNet40Cls(data.Dataset):
def __init__(self,
num_points,
root,
transforms=None,
train=True,
download=True):
def __init__(
self, num_points, root, transforms=None, train=True, download=True
):
super().__init__()
self.transforms = transforms
@@ -37,9 +35,12 @@ class ModelNet40Cls(data.Dataset):
if download and not os.path.exists(self.data_dir):
zipfile = os.path.join(root, os.path.basename(self.url))
subprocess.check_call(
shlex.split("curl {} -o {}".format(self.url, zipfile)))
shlex.split("curl {} -o {}".format(self.url, zipfile))
)
subprocess.check_call(shlex.split("unzip {} -d {}".format(zipfile, root)))
subprocess.check_call(
shlex.split("unzip {} -d {}".format(zipfile, root))
)
subprocess.check_call(shlex.split("rm {}".format(zipfile)))
@@ -83,9 +84,10 @@ class ModelNet40Cls(data.Dataset):
def randomize(self):
self.actual_number_of_points = min(
max(
np.random.randint(self.num_points * 0.8,
self.num_points * 1.2), 1),
self.points.shape[1])
np.random.randint(self.num_points * 0.8, self.num_points * 1.2),
1
), self.points.shape[1]
)
if __name__ == "__main__":
models/Pointnet2Cls.py
+21 -8
View File
@@ -39,6 +39,7 @@ def model_fn_decorator(criterion):
class Pointnet2SSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -48,13 +49,17 @@ class Pointnet2SSG(nn.Module):
npoint=512,
radius=0.2,
nsample=64,
mlp=[input_channels, 64, 64, 128]))
mlp=[input_channels, 64, 64, 128]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=128,
radius=0.4,
nsample=64,
mlp=[128 + 3, 128, 128, 256]))
mlp=[128 + 3, 128, 128, 256]
)
)
self.SA_modules.append(PointnetSAModule(mlp=[256 + 3, 256, 512, 1024]))
self.FC_layer = nn.Sequential(
@@ -62,7 +67,8 @@ class Pointnet2SSG(nn.Module):
nn.Dropout(p=0.5),
pt_utils.FC(512, 256, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(256, num_classes, activation=None))
pt_utils.FC(256, num_classes, activation=None)
)
def forward(self, xyz, points=None):
for module in self.SA_modules:
@@ -72,6 +78,7 @@ class Pointnet2SSG(nn.Module):
class Pointnet2MSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -83,7 +90,9 @@ class Pointnet2MSG(nn.Module):
nsamples=[32, 64, 128],
mlps=[[input_channels, 32, 32,
64], [input_channels, 64, 64, 128],
[input_channels, 64, 96, 128]]))
[input_channels, 64, 96, 128]]
)
)
input_channels = 64 + 128 + 128 + 3
self.SA_modules.append(
@@ -92,17 +101,21 @@ class Pointnet2MSG(nn.Module):
radii=[0.2, 0.4, 0.8],
nsamples=[16, 32, 64],
mlps=[[input_channels, 64, 64,
128], [input_channels, 128, 128, 256],
[input_channels, 128, 128, 256]]))
128], [input_channels, 128, 128, 256],
[input_channels, 128, 128, 256]]
)
)
self.SA_modules.append(
PointnetSAModule(mlp=[128 + 256 + 256 + 3, 256, 512, 1024]))
PointnetSAModule(mlp=[128 + 256 + 256 + 3, 256, 512, 1024])
)
self.FC_layer = nn.Sequential(
pt_utils.FC(1024, 512, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(512, 256, bn=True),
nn.Dropout(p=0.5),
pt_utils.FC(256, num_classes, activation=None))
pt_utils.FC(256, num_classes, activation=None)
)
def forward(self, xyz, points=None):
for module in self.SA_modules:
models/Pointnet2SemSeg.py
+49 -25
View File
@@ -38,6 +38,7 @@ def model_fn_decorator(criterion):
class Pointnet2SSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -49,32 +50,37 @@ class Pointnet2SSG(nn.Module):
npoint=1024,
radius=0.1,
nsample=32,
mlp=[input_channels, 32, 32, 64]))
mlp=[input_channels, 32, 32, 64]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128]))
npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=64,
radius=0.4,
nsample=32,
mlp=[128 + 3, 128, 128, 256]))
npoint=64, radius=0.4, nsample=32, mlp=[128 + 3, 128, 128, 256]
)
)
self.SA_modules.append(
PointnetSAModule(
npoint=16,
radius=0.8,
nsample=32,
mlp=[256 + 3, 256, 256, 512]))
npoint=16, radius=0.8, nsample=32, mlp=[256 + 3, 256, 256, 512]
)
)
self.FP_modules = nn.ModuleList()
self.FP_modules.append(PointnetFPModule(mlp=[128 + input_channels - 3, 128, 128, 128]))
self.FP_modules.append(
PointnetFPModule(mlp=[128 + input_channels - 3, 128, 128, 128])
)
self.FP_modules.append(PointnetFPModule(mlp=[256 + 64, 256, 128]))
self.FP_modules.append(PointnetFPModule(mlp=[256 + 128, 256, 256]))
self.FP_modules.append(PointnetFPModule(mlp=[512 + 256, 256, 256]))
self.FC_layer = nn.Sequential(
pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(),
pt_utils.Conv1d(128, num_classes, activation=None))
pt_utils.Conv1d(128, num_classes, activation=None)
)
def forward(self, xyz, points=None):
if points is not None:
@@ -91,13 +97,17 @@ class Pointnet2SSG(nn.Module):
l_points.append(li_points)
for i in range(-1, -(len(self.FP_modules + 1) - 1), -1):
l_points[i - 1] = self.FP_modules[i](l_xyz[i - 1], l_xyz[i],
l_points[i - 1], l_points[i])
l_points[i - 1] = self.FP_modules[i](
l_xyz[i - 1], l_xyz[i], l_points[i - 1], l_points[i]
)
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
return self.FC_layer(l_points[0].transpose(1,
2)).transpose(1,
2).contiguous()
class Pointnet2MSG(nn.Module):
def __init__(self, num_classes, input_channels=9):
super().__init__()
@@ -111,7 +121,9 @@ class Pointnet2MSG(nn.Module):
npoint=1024,
radii=[0.05, 0.1],
nsamples=[16, 32],
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]))
mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]
)
)
c_out_0 = 32 + 64
c_in = c_out_0 + 3
@@ -120,7 +132,9 @@ class Pointnet2MSG(nn.Module):
npoint=256,
radii=[0.1, 0.2],
nsamples=[16, 32],
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]))
mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]
)
)
c_out_1 = 128 + 128
c_in = c_out_1 + 3
@@ -129,7 +143,9 @@ class Pointnet2MSG(nn.Module):
npoint=64,
radii=[0.2, 0.4],
nsamples=[16, 32],
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]))
mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]
)
)
c_out_2 = 256 + 256
c_in = c_out_2 + 3
@@ -138,20 +154,25 @@ class Pointnet2MSG(nn.Module):
npoint=16,
radii=[0.4, 0.8],
nsamples=[16, 32],
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]))
mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]
)
)
c_out_3 = 512 + 512
self.FP_modules = nn.ModuleList()
self.FP_modules.append(
PointnetFPModule(mlp=[256 + input_channels - 3, 128, 128]))
PointnetFPModule(mlp=[256 + input_channels - 3, 128, 128])
)
self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_0, 256, 256]))
self.FP_modules.append(PointnetFPModule(mlp=[512 + c_out_1, 512, 512]))
self.FP_modules.append(
PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512]))
PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512])
)
self.FC_layer = nn.Sequential(
pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(),
pt_utils.Conv1d(128, num_classes, activation=None))
pt_utils.Conv1d(128, num_classes, activation=None)
)
def forward(self, xyz, points=None):
if points is not None and self.initial_dropout is not None:
@@ -167,10 +188,13 @@ class Pointnet2MSG(nn.Module):
l_points.append(li_points)
for i in range(-1, -(len(self.FP_modules) + 1), -1):
l_points[i - 1] = self.FP_modules[i](l_xyz[i - 1], l_xyz[i],
l_points[i - 1], l_points[i])
l_points[i - 1] = self.FP_modules[i](
l_xyz[i - 1], l_xyz[i], l_points[i - 1], l_points[i]
)
return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous()
return self.FC_layer(l_points[0].transpose(1,
2)).transpose(1,
2).contiguous()
if __name__ == "__main__":
train_cls.py
+46 -31
View File
@@ -16,56 +16,62 @@ import utils.pytorch_utils as pt_utils
import utils.data_utils as d_utils
import argparse
def parse_args():
parser = argparse.ArgumentParser(
description="Arguments for cls training",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
"-batch_size", type=int, default=16, help="Batch size")
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument("-batch_size", type=int, default=16, help="Batch size")
parser.add_argument(
"-num_points",
type=int,
default=1024,
help="Number of points to train with")
help="Number of points to train with"
)
parser.add_argument(
"-weight_decay", type=float, default=1e-5, help="L2 regularization coeff")
parser.add_argument(
"-lr",
"-weight_decay",
type=float,
default=1e-2,
help="Initial learning rate")
default=1e-5,
help="L2 regularization coeff"
)
parser.add_argument(
"-lr_decay",
type=float,
default=0.7,
help="Learning rate decay gamma")
"-lr", type=float, default=1e-2, help="Initial learning rate"
)
parser.add_argument(
"-decay_step",
type=int,
default=20,
help="Learning rate decay step")
"-lr_decay", type=float, default=0.7, help="Learning rate decay gamma"
)
parser.add_argument(
"-decay_step", type=int, default=20, help="Learning rate decay step"
)
parser.add_argument(
"-bn_momentum",
type=float,
default=0.5,
help="Initial batch norm momentum")
help="Initial batch norm momentum"
)
parser.add_argument(
"-bnm_decay",
type=float,
default=0.5,
help="Batch norm momentum decay gamma")
help="Batch norm momentum decay gamma"
)
parser.add_argument(
"-checkpoint", type=str, default=None, help="Checkpoint to start from")
"-checkpoint", type=str, default=None, help="Checkpoint to start from"
)
parser.add_argument(
"-epochs", type=int, default=200, help="Number of epochs to train for")
"-epochs", type=int, default=200, help="Number of epochs to train for"
)
parser.add_argument(
"-run_name",
type=str,
default="cls_run_1",
help="Name for run in tensorboard_logger")
help="Name for run in tensorboard_logger"
)
return parser.parse_args()
lr_clip = 1e-5
bnm_clip = 1e-2
@@ -82,13 +88,15 @@ if __name__ == "__main__":
])
test_set = ModelNet40Cls(
args.num_points, BASE_DIR, transforms=transforms, train=False)
args.num_points, BASE_DIR, transforms=transforms, train=False
)
test_loader = DataLoader(
test_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
pin_memory=True)
pin_memory=True
)
train_set = ModelNet40Cls(args.num_points, BASE_DIR, transforms=transforms)
train_loader = DataLoader(
@@ -96,25 +104,30 @@ if __name__ == "__main__":
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
pin_memory=True)
pin_memory=True
)
tb_log.configure('runs/{}'.format(args.run_name))
model = Pointnet(input_channels=3, num_classes=40)
model.cuda()
optimizer = optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
model.parameters(), lr=args.lr, weight_decay=args.weight_decay
)
lr_lbmd = lambda e: max(args.lr_decay**(e // args.decay_step), lr_clip / args.lr)
bn_lbmd = lambda e: max(args.bn_momentum * args.bnm_decay**(e // args.decay_step), bnm_clip)
if args.checkpoint is not None:
start_epoch, best_loss = pt_utils.load_checkpoint(
model, optimizer, filename=args.checkpoint.split(".")[0])
model, optimizer, filename=args.checkpoint.split(".")[0]
)
lr_scheduler = lr_sched.LambdaLR(
optimizer, lr_lambda=lr_lbmd, last_epoch=start_epoch)
optimizer, lr_lambda=lr_lbmd, last_epoch=start_epoch
)
bnm_scheduler = pt_utils.BNMomentumScheduler(
model, bn_lambda=bn_lbmd, last_epoch=start_epoch)
model, bn_lambda=bn_lbmd, last_epoch=start_epoch
)
else:
lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lambda=lr_lbmd)
bnm_scheduler = pt_utils.BNMomentumScheduler(model, bn_lambda=bn_lbmd)
@@ -131,14 +144,16 @@ if __name__ == "__main__":
checkpoint_name="cls_checkpoint",
best_name="cls_best",
lr_scheduler=lr_scheduler,
bnm_scheduler=bnm_scheduler)
bnm_scheduler=bnm_scheduler
)
trainer.train(
start_epoch,
args.epochs,
train_loader,
test_loader,
best_loss=best_loss)
best_loss=best_loss
)
if start_epoch == args.epochs:
_ = trainer.eval_epoch(start_epoch, test_loader)
train_sem_seg.py
+40 -20
View File
@@ -18,51 +18,62 @@ import argparse
parser = argparse.ArgumentParser(description="Arg parser")
parser.add_argument(
"-batch_size", type=int, default=32, help="Batch size [default: 32]")
"-batch_size", type=int, default=32, help="Batch size [default: 32]"
)
parser.add_argument(
"-num_points",
type=int,
default=2048,
help="Number of points to train with [default: 2048]")
help="Number of points to train with [default: 2048]"
)
parser.add_argument(
"-weight_decay",
type=float,
default=0,
help="L2 regularization coeff [default: 0.0]")
help="L2 regularization coeff [default: 0.0]"
)
parser.add_argument(
"-lr",
type=float,
default=1e-2,
help="Initial learning rate [default: 1e-2]")
help="Initial learning rate [default: 1e-2]"
)
parser.add_argument(
"-lr_decay",
type=float,
default=0.5,
help="Learning rate decay gamma [default: 0.5]")
help="Learning rate decay gamma [default: 0.5]"
)
parser.add_argument(
"-decay_step",
type=int,
default=20,
help="Learning rate decay step [default: 20]")
help="Learning rate decay step [default: 20]"
)
parser.add_argument(
"-bn_momentum",
type=float,
default=0.9,
help="Initial batch norm momentum [default: 0.9]")
help="Initial batch norm momentum [default: 0.9]"
)
parser.add_argument(
"-bn_decay",
type=float,
default=0.5,
help="Batch norm momentum decay gamma [default: 0.5]")
help="Batch norm momentum decay gamma [default: 0.5]"
)
parser.add_argument(
"-checkpoint", type=str, default=None, help="Checkpoint to start from")
"-checkpoint", type=str, default=None, help="Checkpoint to start from"
)
parser.add_argument(
"-epochs", type=int, default=200, help="Number of epochs to train for")
"-epochs", type=int, default=200, help="Number of epochs to train for"
)
parser.add_argument(
"-run_name",
type=str,
default="sem_seg_run_1",
help="Name for run in tensorboard_logger")
help="Name for run in tensorboard_logger"
)
BASE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
@@ -74,13 +85,15 @@ if __name__ == "__main__":
tb_log.configure('runs/{}'.format(args.run_name))
test_set = Indoor3DSemSeg(
args.num_points, BASE_DIR, train=False, data_precent=0.01)
args.num_points, BASE_DIR, train=False, data_precent=0.01
)
test_loader = DataLoader(
test_set,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=2)
num_workers=2
)
train_set = Indoor3DSemSeg(args.num_points, BASE_DIR, data_precent=1.0)
train_loader = DataLoader(
@@ -88,12 +101,14 @@ if __name__ == "__main__":
batch_size=args.batch_size,
pin_memory=True,
num_workers=2,
shuffle=True)
shuffle=True
)
model = Pointnet(num_classes=13)
model.cuda()
optimizer = optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
model.parameters(), lr=args.lr, weight_decay=args.weight_decay
)
lr_lbmd = lambda e: max(args.lr_decay**(e // args.decay_step), lr_clip / args.lr)
bnm_lmbd = lambda e: max(args.bn_momentum * args.bn_decay**(e // args.decay_step), bnm_clip)
@@ -106,12 +121,15 @@ if __name__ == "__main__":
best_loss = 1e10
else:
start_epoch, best_loss = pt_utils.load_checkpoint(
model, optimizer, filename=args.checkpoint.split(".")[0])
model, optimizer, filename=args.checkpoint.split(".")[0]
)
lr_scheduler = lr_sched.LambdaLR(
optimizer, lr_lbmd, last_epoch=start_epoch)
optimizer, lr_lbmd, last_epoch=start_epoch
)
bnm_scheduler = pt_utils.BNMomentumScheduler(
model, bnm_lmbd, last_epoch=start_epoch)
model, bnm_lmbd, last_epoch=start_epoch
)
model_fn = model_fn_decorator(nn.CrossEntropyLoss())
@@ -123,14 +141,16 @@ if __name__ == "__main__":
best_name="sem_seg_best",
lr_scheduler=lr_scheduler,
bnm_scheduler=bnm_scheduler,
eval_frequency=10)
eval_frequency=10
)
trainer.train(
start_epoch,
args.epochs,
train_loader,
test_loader,
best_loss=best_loss)
best_loss=best_loss
)
if start_epoch == args.epochs:
test_loader.dataset.data_precent = 1.0
utils/build_ffi.py
+7 -3
View File
@@ -9,7 +9,8 @@ base_dir = osp.dirname(osp.abspath(__file__))
def parse_args():
parser = argparse.ArgumentParser(
description="Arguments for building pointnet2 ffi extension")
description="Arguments for building pointnet2 ffi extension"
)
parser.add_argument("--objs", nargs="*")
clean_arg = parser.add_mutually_exclusive_group()
clean_arg.add_argument("--build", dest='build', action="store_true")
@@ -27,7 +28,7 @@ def build(args):
extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')]
ffi = create_extension(
'_ext',
'_ext.pointnet2',
headers=[a for a in glob.glob("cinclude/*_wrapper.h")],
sources=[a for a in glob.glob("csrc/*.c")],
define_macros=[('WITH_CUDA', None)],
@@ -36,12 +37,15 @@ def build(args):
extra_objects=extra_objects,
include_dirs=[osp.join(base_dir, 'cinclude')],
verbose=False,
package=False)
package=False
)
ffi.build()
def clean(args):
shutil.rmtree(osp.join(base_dir, "_ext"))
if __name__ == "__main__":
args = parse_args()
if args.clean:
utils/csrc/ball_query.c
+7 -7
View File
@@ -8,13 +8,13 @@ int ball_query_wrapper(int b, int n, int m, float radius, int nsample,
THCudaTensor *new_xyz_tensor, THCudaTensor *xyz_tensor,
THCudaIntTensor *idx_tensor) {
const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor);
const float *xyz = THCudaTensor_data(state, xyz_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor);
const float *xyz = THCudaTensor_data(state, xyz_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz,
idx, stream);
return 1;
query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz, idx,
stream);
return 1;
}
utils/csrc/ball_query_gpu.cu
+35 -37
View File
@@ -11,38 +11,37 @@ __global__ void query_ball_point_kernel(int b, int n, int m, float radius,
int nsample,
const float *__restrict__ new_xyz,
const float *__restrict__ xyz,
int * __restrict__ idx) {
int batch_index = blockIdx.x;
xyz += batch_index * n * 3;
new_xyz += batch_index * m * 3;
idx += m * nsample * batch_index;
int *__restrict__ idx) {
int batch_index = blockIdx.x;
xyz += batch_index * n * 3;
new_xyz += batch_index * m * 3;
idx += m * nsample * batch_index;
int index = threadIdx.x;
int stride = blockDim.x;
int index = threadIdx.x;
int stride = blockDim.x;
float radius2 = radius * radius;
for (int j = index; j < m; j += stride) {
float new_x = new_xyz[j * 3 + 0];
float new_y = new_xyz[j * 3 + 1];
float new_z = new_xyz[j * 3 + 2];
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (new_x - x) * (new_x - x) +
(new_y - y) * (new_y - y) +
(new_z - z) * (new_z - z);
if (d2 < radius2) {
if (cnt == 0) {
for (int l = 0; l < nsample; ++l) {
idx[j * nsample + l] = k;
}
}
idx[j * nsample + cnt] = k;
++cnt;
}
float radius2 = radius * radius;
for (int j = index; j < m; j += stride) {
float new_x = new_xyz[j * 3 + 0];
float new_y = new_xyz[j * 3 + 1];
float new_z = new_xyz[j * 3 + 2];
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
(new_z - z) * (new_z - z);
if (d2 < radius2) {
if (cnt == 0) {
for (int l = 0; l < nsample; ++l) {
idx[j * nsample + l] = k;
}
}
idx[j * nsample + cnt] = k;
++cnt;
}
}
}
}
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
@@ -50,14 +49,13 @@ void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
const float *xyz, int *idx,
cudaStream_t stream) {
cudaError_t err;
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
b, n, m, radius, nsample, new_xyz, xyz, idx);
cudaError_t err;
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
b, n, m, radius, nsample, new_xyz, xyz, idx);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
utils/csrc/group_points.c
+14 -14
View File
@@ -9,15 +9,15 @@ int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
THCudaIntTensor *idx_tensor,
THCudaTensor *out_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
stream);
return 1;
group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
stream);
return 1;
}
int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
@@ -25,13 +25,13 @@ int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
THCudaIntTensor *idx_tensor,
THCudaTensor *grad_points_tensor) {
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out,
idx, grad_points, stream);
return 1;
group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out, idx,
grad_points, stream);
return 1;
}
utils/csrc/group_points_gpu.cu
+49 -52
View File
@@ -1,86 +1,83 @@
#include <stdio.h>
#include <stdlib.h>
#include "group_points_gpu.h"
#include "cuda_utils.h"
#include "group_points_gpu.h"
// input: points(b, n, c) idx(b, npoints, nsample)
// output: out(b, npoints, nsample, c)
__global__ void group_points_kernel(int b, int n, int c, int npoints,
int nsample,
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * n * c;
idx += batch_index * npoints * nsample;
out += batch_index * npoints * nsample * c;
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * n * c;
idx += batch_index * npoints * nsample;
out += batch_index * npoints * nsample * c;
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
memcpy(out + j * nsample * c + k * c, points + ii * c,
sizeof(float) * c);
}
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
memcpy(out + j * nsample * c + k * c, points + ii * c,
sizeof(float) * c);
}
}
}
void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
const float *points, const int *idx,
float *out, cudaStream_t stream) {
cudaError_t err;
group_points_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, points, idx, out);
cudaError_t err;
group_points_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, points, idx, out);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
// input: grad_out(b, npoints, nsample, c), idx(b, npoints, nsample)
// output: grad_points(b, n, c)
__global__ void group_points_grad_kernel(int b, int n, int c, int npoints,
int nsample,
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_points += batch_index * n * c;
idx += batch_index * npoints * nsample;
grad_out += batch_index * npoints * nsample * c;
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_points += batch_index * n * c;
idx += batch_index * npoints * nsample;
grad_out += batch_index * npoints * nsample * c;
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
for (int l = 0; l < c; ++l) {
atomicAdd(
grad_points + ii * c + l,
grad_out[j * nsample * c + k * c + l]);
}
}
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
for (int l = 0; l < c; ++l) {
atomicAdd(grad_points + ii * c + l,
grad_out[j * nsample * c + k * c + l]);
}
}
}
}
void group_points_grad_kernel_wrapper(int b, int n, int c, int npoints,
int nsample, const float *grad_out,
const int *idx, float *grad_points,
cudaStream_t stream) {
cudaError_t err;
group_points_grad_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, grad_out, idx, grad_points);
cudaError_t err;
group_points_grad_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, grad_out, idx, grad_points);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
utils/csrc/interpolate_gpu.cu
+1 -1
View File
@@ -2,8 +2,8 @@
#include <stdio.h>
#include <stdlib.h>
#include "interpolate_gpu.h"
#include "cuda_utils.h"
#include "interpolate_gpu.h"
// input: unknown(b, n, 3) known(b, m, 3)
// output: dist2(b, n, 3), idx(b, n, 3)
utils/csrc/roi_mask_points_wrapper.c
+28 -29
View File
@@ -9,17 +9,17 @@ int roi_mask_wrapper(int n_roi, int b, int n, THCudaTensor *rois_tensor,
THCudaTensor *data_xyz_tensor,
THCudaByteTensor *mask_tensor) {
const float *rois = THCudaTensor_data(state, rois_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const float *data_xyz = THCudaTensor_data(state, data_xyz_tensor);
unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *rois = THCudaTensor_data(state, rois_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const float *data_xyz = THCudaTensor_data(state, data_xyz_tensor);
unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_mask_kernel_wrapper(n_roi, b, n, rois, batch_indices, data_xyz,
mask, stream);
return 1;
roi_mask_kernel_wrapper(n_roi, b, n, rois, batch_indices, data_xyz, mask,
stream);
return 1;
}
int roi_avg_pool_forward_wrapper(int n_roi, int b, int n, int d,
@@ -28,17 +28,17 @@ int roi_avg_pool_forward_wrapper(int n_roi, int b, int n, int d,
THCudaTensor *points_tensor,
THCudaTensor *descriptors_tensor) {
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *descriptors = THCudaTensor_data(state, descriptors_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *descriptors = THCudaTensor_data(state, descriptors_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_forward_wrapper(n_roi, b, n, d, mask, batch_indices,
points, descriptors, stream);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_forward_wrapper(n_roi, b, n, d, mask, batch_indices,
points, descriptors, stream);
return 1;
return 1;
}
int roi_avg_pool_backward_wrapper(int n_roi, int b, int n, int d,
@@ -47,17 +47,16 @@ int roi_avg_pool_backward_wrapper(int n_roi, int b, int n, int d,
THCudaTensor *grad_descriptors_tensor,
THCudaTensor *grad_points_tensor) {
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *grad_descriptors =
THCudaTensor_data(state, grad_descriptors_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *grad_descriptors =
THCudaTensor_data(state, grad_descriptors_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_backward_wrapper(n_roi, b, n, d, mask,
batch_indices, grad_descriptors,
grad_points, stream);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_backward_wrapper(n_roi, b, n, d, mask, batch_indices,
grad_descriptors, grad_points, stream);
return 1;
return 1;
}
utils/csrc/sampling.c
+12 -14
View File
@@ -9,15 +9,14 @@ int gather_points_wrapper(int b, int n, int c, int npoints,
THCudaIntTensor *idx_tensor,
THCudaTensor *out_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out,
stream);
return 1;
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out, stream);
return 1;
}
int furthest_point_sampling_wrapper(int b, int n, int m,
@@ -25,13 +24,12 @@ int furthest_point_sampling_wrapper(int b, int n, int m,
THCudaTensor *temp_tensor,
THCudaIntTensor *idx_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
float *temp = THCudaTensor_data(state, temp_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *temp = THCudaTensor_data(state, temp_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx,
stream);
return 1;
furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx, stream);
return 1;
}
utils/data_utils.py
+11 -4
View File
@@ -3,17 +3,19 @@ import numpy as np
class PointcloudScale(object):
def __init__(self, mean=2.0, std=1.0, clip=1.8):
self.mean, self.std, self.clip = mean, std, clip
def __call__(self, points):
scaler = points.new(1).normal_(
mean=self.mean, std=self.std).clamp_(
max(self.mean - self.clip, 0.01), self.mean + self.clip)
mean=self.mean, std=self.std
).clamp_(max(self.mean - self.clip, 0.01), self.mean + self.clip)
return scaler * points
class PointcloudRotate(object):
def __init__(self, x_axis=False, z_axis=True):
assert x_axis or z_axis
self.x, self.z = x_axis, z_axis
@@ -46,25 +48,30 @@ class PointcloudRotate(object):
class PointcloudJitter(object):
def __init__(self, std=0.01, clip=0.03):
self.std, self.clip = std, clip
def __call__(self, points):
jittered_data = points.new(*points.size()).normal_(
mean=0.0, std=self.std).clamp_(-self.clip, self.clip)
mean=0.0, std=self.std
).clamp_(-self.clip, self.clip)
return points + jittered_data
class PointcloudTranslate(object):
def __init__(self, std=1.0, clip=3.0):
self.std, self.clip = std, clip
def __call__(self, points):
translation = points.new(3).normal_(
mean=0.0, std=self.std).clamp_(-self.clip, self.clip)
mean=0.0, std=self.std
).clamp_(-self.clip, self.clip)
return points + translation
class PointcloudToTensor(object):
def __call__(self, points):
return torch.from_numpy(points).float()
utils/linalg_utils.py
+5 -3
View File
@@ -4,9 +4,11 @@ from enum import Enum
PDist2Order = Enum('PDist2Order', 'd_first d_second')
def pdist2(X: torch.Tensor,
Z: torch.Tensor = None,
order: PDist2Order = PDist2Order.d_second) -> torch.Tensor:
def pdist2(
X: torch.Tensor,
Z: torch.Tensor = None,
order: PDist2Order = PDist2Order.d_second
) -> torch.Tensor:
r""" Calculates the pairwise distance between X and Z
D[b, i, j] = l2 distance X[b, i] and Z[b, j]
utils/pointnet2_modules.py
+54 -44
View File
@@ -24,13 +24,15 @@ class PointnetSAModuleMSG(nn.Module):
Use batchnorm
"""
def __init__(self,
*,
npoint: int,
radii: List[float],
nsamples: List[int],
mlps: List[List[int]],
bn: bool = True):
def __init__(
self,
*,
npoint: int,
radii: List[float],
nsamples: List[int],
mlps: List[List[int]],
bn: bool = True
):
super().__init__()
assert len(radii) == len(nsamples) == len(mlps)
@@ -41,8 +43,7 @@ class PointnetSAModuleMSG(nn.Module):
for i in range(len(radii)):
radius = radii[i]
nsample = nsamples[i]
self.groupers.append(
pointnet2_utils.QueryAndGroup(radius, nsample))
self.groupers.append(pointnet2_utils.QueryAndGroup(radius, nsample))
mlp_spec = mlps[i]
self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn))
@@ -66,18 +67,20 @@ class PointnetSAModuleMSG(nn.Module):
new_points_list = []
new_xyz = pointnet2_utils.gather_points(
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint))
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
)
for i in range(len(self.groupers)):
new_points = self.groupers[i](xyz, new_xyz, points)
new_points = self.mlps[i](new_points.permute(
0, 3, 1, 2)) # (B, mlp[-1], npoint, nsample)
new_points = self.mlps[i](new_points.permute(0, 3, 1, 2)
) # (B, mlp[-1], npoint, nsample)
new_points = F.max_pool2d(
new_points,
kernel_size=[1, new_points.size(3)]) # (B, mlp[-1], npoint, 1)
new_points, kernel_size=[1, new_points.size(3)]
) # (B, mlp[-1], npoint, 1)
new_points = new_points.squeeze(-1) # (B, mlp[-1], npoint)
new_points = new_points.transpose(
1, 2).contiguous() # (B, npoint, mlp[-1])
1, 2
).contiguous() # (B, npoint, mlp[-1])
new_points_list.append(new_points)
@@ -101,13 +104,15 @@ class PointnetSAModule(nn.Module):
Use batchnorm
"""
def __init__(self,
*,
mlp: List[int],
npoint: int = None,
radius: float = None,
nsample: int = None,
bn: bool = True):
def __init__(
self,
*,
mlp: List[int],
npoint: int = None,
radius: float = None,
nsample: int = None,
bn: bool = True
):
super().__init__()
self.npoint = npoint
@@ -140,21 +145,23 @@ class PointnetSAModule(nn.Module):
if self.npoint is not None:
new_xyz = pointnet2_utils.gather_points(
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint))
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
)
else:
new_xyz = xyz.data.new([[[0, 0, 0]]]).expand(xyz.size(0), 1, 3)
new_points = self.grouper(xyz, new_xyz,
points) # (B, npoint, nsample, 3 + C)
new_points = self.grouper(
xyz, new_xyz, points
) # (B, npoint, nsample, 3 + C)
new_points = self.mlp(new_points.permute(
0, 3, 1, 2)) # (B, mlp[-1], npoint, nsample)
new_points = self.mlp(new_points.permute(0, 3, 1, 2)
) # (B, mlp[-1], npoint, nsample)
new_points = F.max_pool2d(
new_points,
kernel_size=[1, new_points.size(3)]) # (B, mlp[-1], npoint, 1)
new_points, kernel_size=[1, new_points.size(3)]
) # (B, mlp[-1], npoint, 1)
new_points = new_points.squeeze(-1) # (B, mlp[-1], npoint)
new_points = new_points.transpose(
1, 2).contiguous() # (B, npoint, mlp[-1])
new_points = new_points.transpose(1, 2
).contiguous() # (B, npoint, mlp[-1])
return new_xyz, new_points
@@ -174,9 +181,10 @@ class PointnetFPModule(nn.Module):
super().__init__()
self.mlp = pt_utils.SharedMLP(mlp, bn=bn)
def forward(self, unknown: torch.Tensor, known: torch.Tensor,
unknow_feats: torch.Tensor,
known_feats: torch.Tensor) -> torch.Tensor:
def forward(
self, unknown: torch.Tensor, known: torch.Tensor,
unknow_feats: torch.Tensor, known_feats: torch.Tensor
) -> torch.Tensor:
r"""
Parameters
----------
@@ -201,19 +209,21 @@ class PointnetFPModule(nn.Module):
weight = dist_recip / norm
interpolated_feats = pointnet2_utils.three_interpolate(
known_feats, idx, weight)
known_feats, idx, weight
)
if unknow_feats is not None:
new_points = torch.cat(
[interpolated_feats, unknow_feats], dim=-1) #(B, n, C2 + C1)
new_points = torch.cat([interpolated_feats, unknow_feats],
dim=-1) #(B, n, C2 + C1)
else:
new_points = interpolated_feats
new_points = new_points.unsqueeze(-1).transpose(1,
2) #(B, C2 + C1, n, 1)
new_points = new_points.unsqueeze(-1).transpose(
1, 2
) #(B, C2 + C1, n, 1)
new_points = self.mlp(new_points)
return new_points.squeeze(-1).transpose(
1, 2).contiguous() #(B, n, mlp[-1])
return new_points.squeeze(-1).transpose(1, 2
).contiguous() #(B, n, mlp[-1])
if __name__ == "__main__":
@@ -224,7 +234,8 @@ if __name__ == "__main__":
xyz_feats = Variable(torch.randn(2, 10, 6).cuda(), requires_grad=True)
test_module = PointnetSAModuleMSG(
npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]])
npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]]
)
test_module.cuda()
print(test_module(xyz, xyz_feats))
@@ -237,7 +248,6 @@ if __name__ == "__main__":
for _ in range(1):
_, new_points = test_module(xyz, xyz_feats)
new_points.backward(
torch.cuda.FloatTensor(*new_points.size()).fill_(1))
new_points.backward(torch.cuda.FloatTensor(*new_points.size()).fill_(1))
print(new_points)
print(xyz.grad)
utils/pointnet2_utils.py
+57 -36
View File
@@ -5,12 +5,14 @@ import torch.nn.functional as F
import torch.nn as nn
from linalg_utils import pdist2, PDist2Order
from collections import namedtuple
import _ext as pointnet2
import pytorch_utils as pt_utils
from typing import List, Tuple
from _ext import pointnet2
class RandomDropout(nn.Module):
def __init__(self, p=0.5, inplace=False):
super().__init__()
self.p = p
@@ -18,11 +20,13 @@ class RandomDropout(nn.Module):
def forward(self, X):
theta = torch.Tensor(1).uniform_(0, self.p)[0]
return pt_utils.feature_dropout_no_scaling(X, theta, self.train,
self.inplace)
return pt_utils.feature_dropout_no_scaling(
X, theta, self.train, self.inplace
)
class FurthestPointSampling(Function):
@staticmethod
def forward(ctx, xyz: torch.Tensor, npoint: int) -> torch.Tensor:
r"""
@@ -30,16 +34,16 @@ class FurthestPointSampling(Function):
minimum distance
Parameters
---------
----------
xyz : torch.Tensor
(B, N, 3) tensor where N > npoint
npoint : int32
number of points in the sampled set
Returns
-------
torch.Tensor
(B, npoint) tensor containing the set
------
"""
B, N, _ = xyz.size()
@@ -50,8 +54,9 @@ class FurthestPointSampling(Function):
temp = temp.contiguous()
output = output.contiguous()
pointnet2.furthest_point_sampling_wrapper(B, N, npoint, xyz, temp,
output)
pointnet2.furthest_point_sampling_wrapper(
B, N, npoint, xyz, temp, output
)
return output
@@ -64,6 +69,7 @@ furthest_point_sample = FurthestPointSampling.apply
class GatherPoints(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
r"""
@@ -71,7 +77,7 @@ class GatherPoints(Function):
minimum distance
Parameters
---------
----------
points : torch.Tensor
(B, N, 3) tensor
@@ -79,9 +85,9 @@ class GatherPoints(Function):
(B, npoint) tensor of the points to gather
Returns
-------
torch.Tensor
(B, npoint, 3) tensor
------
"""
B, N, C = points.size()
@@ -106,6 +112,7 @@ gather_points = GatherPoints.apply
class ThreeNN(Function):
@staticmethod
def forward(ctx, unknown: torch.Tensor,
known: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -147,9 +154,11 @@ three_nn = ThreeNN.apply
class ThreeInterpolate(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor,
weight: torch.Tensor) -> torch.Tensor:
def forward(
ctx, points: torch.Tensor, idx: torch.Tensor, weight: torch.Tensor
) -> torch.Tensor:
r"""
Performs weight linear interpolation on 3 points
Parameters
@@ -178,14 +187,15 @@ class ThreeInterpolate(Function):
idx = idx.contiguous()
weight = weight.contiguous()
output = output.contiguous()
pointnet2.three_interpolate_wrapper(B, m, c, n, points, idx, weight,
output)
pointnet2.three_interpolate_wrapper(
B, m, c, n, points, idx, weight, output
)
return output
@staticmethod
def backward(ctx, grad_out: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
r"""
Parameters
----------
@@ -196,6 +206,7 @@ class ThreeInterpolate(Function):
-------
grad_points : torch.Tensor
(B, m, c) tensor with gradients of points
None
None
@@ -209,8 +220,9 @@ class ThreeInterpolate(Function):
idx = idx.contiguous()
weight = weight.contiguous()
grad_points = grad_points.contiguous()
pointnet2.three_interpolate_grad_wrapper(B, n, c, m, grad_out.data,
idx, weight, grad_points.data)
pointnet2.three_interpolate_grad_wrapper(
B, n, c, m, grad_out.data, idx, weight, grad_points.data
)
return grad_points, None, None
@@ -219,6 +231,7 @@ three_interpolate = ThreeInterpolate.apply
class GroupPoints(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
r"""
@@ -243,8 +256,9 @@ class GroupPoints(Function):
points = points.contiguous()
idx = idx.contiguous()
output = output.contiguous()
pointnet2.group_points_wrapper(B, N, C, npoints, nsample, points, idx,
output)
pointnet2.group_points_wrapper(
B, N, C, npoints, nsample, points, idx, output
)
ctx.idx_N_C_for_backward = (idx, N, C)
return output
@@ -273,7 +287,8 @@ class GroupPoints(Function):
grad_out = grad_out.contiguous()
grad_points = grad_points.contiguous()
pointnet2.group_points_grad_wrapper(
B, N, C, npoint, nsample, grad_out.data, idx, grad_points.data)
B, N, C, npoint, nsample, grad_out.data, idx, grad_points.data
)
return grad_points, None
@@ -282,13 +297,16 @@ group_points = GroupPoints.apply
class BallQuery(Function):
@staticmethod
def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor,
new_xyz: torch.Tensor) -> torch.Tensor:
def forward(
ctx, radius: float, nsample: int, xyz: torch.Tensor,
new_xyz: torch.Tensor
) -> torch.Tensor:
r"""
Parameters
---------
----------
radius : float
radius of the balls
nsample : int
@@ -299,7 +317,7 @@ class BallQuery(Function):
(B, npoint, 3) centers of the ball query
Returns
------
-------
torch.Tensor
(B, npoint, nsample) tensor with the indicies of the points that form the query balls
"""
@@ -311,8 +329,9 @@ class BallQuery(Function):
new_xyz = new_xyz.contiguous()
xyz = xyz.contiguous()
idx = idx.contiguous()
pointnet2.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz,
xyz, idx)
pointnet2.ball_query_wrapper(
B, N, npoint, radius, nsample, new_xyz, xyz, idx
)
return idx
@@ -344,10 +363,11 @@ class QueryAndGroup(nn.Module):
self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
points: torch.Tensor = None) -> Tuple[torch.Tensor]:
points: torch.Tensor = None
) -> Tuple[torch.Tensor]:
r"""
Parameters
---------
----------
xyz : torch.Tensor
xyz coordinates of the points (B, N, 3)
new_xyz : torch.Tensor
@@ -368,9 +388,8 @@ class QueryAndGroup(nn.Module):
if points is not None:
grouped_points = group_points(points, idx)
if self.use_xyz:
new_points = torch.cat(
[grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
new_points = torch.cat([grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
else:
new_points = group_points
else:
@@ -395,10 +414,11 @@ class GroupAll(nn.Module):
self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
points: torch.Tensor = None) -> Tuple[torch.Tensor]:
points: torch.Tensor = None
) -> Tuple[torch.Tensor]:
r"""
Parameters
---------
----------
xyz : torch.Tensor
xyz coordinates of the points (B, N, 3)
new_xyz : torch.Tensor
@@ -414,11 +434,12 @@ class GroupAll(nn.Module):
grouped_xyz = xyz.view(xyz.size(0), 1, xyz.size(1), xyz.size(2))
if points is not None:
grouped_points = points.view(points.size(0), 1, points.size(1), points.size(2))
grouped_points = points.view(
points.size(0), 1, points.size(1), points.size(2)
)
if self.use_xyz:
new_points = torch.cat(
[grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
new_points = torch.cat([grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
else:
new_points = group_points
else:
utils/pytorch_utils.py
+182 -141
View File
@@ -16,48 +16,55 @@ import math
class SharedMLP(nn.Sequential):
def __init__(self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""):
def __init__(
self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""
):
super().__init__()
for i in range(len(args) - 1):
self.add_module(name + 'layer{}'.format(i),
Conv2d(
args[i],
args[i + 1],
bn=bn,
activation=activation))
self.add_module(
name + 'layer{}'.format(i),
Conv2d(args[i], args[i + 1], bn=bn, activation=activation)
)
class _ConvBase(nn.Sequential):
def __init__(self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""):
def __init__(
self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""
):
super().__init__()
bias = bias and (not bn)
self.add_module(name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias))
self.add_module(
name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias
)
)
init(self[0].weight)
if bias:
@@ -73,18 +80,21 @@ class _ConvBase(nn.Sequential):
class Conv1d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -97,22 +107,26 @@ class Conv1d(_ConvBase):
conv=nn.Conv1d,
batch_norm=nn.BatchNorm1d,
bias=bias,
name=name)
name=name
)
class Conv2d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -125,22 +139,26 @@ class Conv2d(_ConvBase):
conv=nn.Conv2d,
batch_norm=nn.BatchNorm2d,
bias=bias,
name=name)
name=name
)
class Conv3d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -153,18 +171,22 @@ class Conv3d(_ConvBase):
conv=nn.Conv3d,
batch_norm=nn.BatchNorm3d,
bias=bias,
name=name)
name=name
)
class FC(nn.Sequential):
def __init__(self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""
):
super().__init__()
self.add_module(name + 'fc', nn.Linear(in_size, out_size, bias=not bn))
if init is not None:
@@ -183,6 +205,7 @@ class FC(nn.Sequential):
class _DropoutNoScaling(InplaceFunction):
@staticmethod
def _make_noise(input):
return input.new().resize_as_(input)
@@ -192,8 +215,9 @@ class _DropoutNoScaling(InplaceFunction):
if inplace:
return None
n = g.appendNode(
g.create("Dropout", [input]).f_("ratio", p).i_(
"is_test", not train))
g.create("Dropout", [input]).f_("ratio",
p).i_("is_test", not train)
)
real = g.appendNode(g.createSelect(n, 0))
g.appendNode(g.createSelect(n, 1))
return real
@@ -201,8 +225,10 @@ class _DropoutNoScaling(InplaceFunction):
@classmethod
def forward(cls, ctx, input, p=0.5, train=False, inplace=False):
if p < 0 or p > 1:
raise ValueError("dropout probability has to be between 0 and 1, "
"but got {}".format(p))
raise ValueError(
"dropout probability has to be between 0 and 1, "
"but got {}".format(p)
)
ctx.p = p
ctx.train = train
ctx.inplace = inplace
@@ -236,6 +262,7 @@ dropout_no_scaling = _DropoutNoScaling.apply
class _FeatureDropoutNoScaling(_DropoutNoScaling):
@staticmethod
def symbolic(input, p=0.5, train=False, inplace=False):
return None
@@ -244,7 +271,8 @@ class _FeatureDropoutNoScaling(_DropoutNoScaling):
def _make_noise(input):
return input.new().resize_(
input.size(0), input.size(1), *repeat(1,
input.dim() - 2))
input.dim() - 2)
)
feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
@@ -252,21 +280,17 @@ feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
def checkpoint_state(model=None, optimizer=None, best_prec=None, epoch=None):
return {
'epoch':
epoch,
'best_prec':
best_prec,
'model_state':
model.state_dict() if model is not None else None,
'optimizer_state':
optimizer.state_dict() if optimizer is not None else None
'epoch': epoch,
'best_prec': best_prec,
'model_state': model.state_dict() if model is not None else None,
'optimizer_state': optimizer.state_dict()
if optimizer is not None else None
}
def save_checkpoint(state,
is_best,
filename='checkpoint',
bestname='model_best'):
def save_checkpoint(
state, is_best, filename='checkpoint', bestname='model_best'
):
filename = '{}.pth.tar'.format(filename)
torch.save(state, filename)
if is_best:
@@ -325,7 +349,8 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
out = out.view(
len(batch), max_len,
*[batch[0].size(i) for i in range(1, batch[0].dim())])
*[batch[0].size(i) for i in range(1, batch[0].dim())]
)
out.fill_(pad_val)
for i in range(len(batch)):
out[i, 0:batch[i].size(0)] = batch[i]
@@ -342,8 +367,9 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
return wrapped([torch.from_numpy(b) for b in batch])
if elem.shape == (): # scalars
py_type = float if elem.dtype.name.startswith('float') else int
return _numpy_type_map[elem.dtype.name](list(
map(py_type, batch)))
return _numpy_type_map[elem.dtype.name](
list(map(py_type, batch))
)
elif isinstance(batch[0], int):
return torch.LongTensor(batch)
elif isinstance(batch[0], float):
@@ -372,19 +398,19 @@ class TrainValSplitter():
Whether or not shuffle which data goes to which split
"""
def __init__(self,
*,
numel: int,
percent_train: float,
shuffled: bool = False):
def __init__(
self, *, numel: int, percent_train: float, shuffled: bool = False
):
indicies = np.array([i for i in range(numel)])
if shuffled:
np.random.shuffle(indicies)
self.train = torch.utils.data.sampler.SubsetRandomSampler(
indicies[0:int(percent_train * numel)])
indicies[0:int(percent_train * numel)]
)
self.val = torch.utils.data.sampler.SubsetRandomSampler(
indicies[int(percent_train * numel):-1])
indicies[int(percent_train * numel):-1]
)
class CrossValSplitter():
@@ -413,7 +439,8 @@ class CrossValSplitter():
self.val = torch.utils.data.sampler.SubsetRandomSampler(self.folds[0])
self.train = torch.utils.data.sampler.SubsetRandomSampler(
np.concatenate(self.folds[1:], axis=0))
np.concatenate(self.folds[1:], axis=0)
)
self.metrics = {}
@@ -428,7 +455,8 @@ class CrossValSplitter():
assert idx >= 0 and idx < len(self)
self.val.inidicies = self.folds[idx]
self.train.inidicies = np.concatenate(
self.folds[np.arange(len(self)) != idx], axis=0)
self.folds[np.arange(len(self)) != idx], axis=0
)
def __next__(self):
self.current_v_ind += 1
@@ -454,6 +482,7 @@ class CrossValSplitter():
def set_bn_momentum_default(bn_momentum):
def fn(m):
if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
m.momentum = bn_momentum
@@ -462,14 +491,17 @@ def set_bn_momentum_default(bn_momentum):
class BNMomentumScheduler(object):
def __init__(self,
model,
bn_lambda,
last_epoch=-1,
setter=set_bn_momentum_default):
def __init__(
self, model, bn_lambda, last_epoch=-1,
setter=set_bn_momentum_default
):
if not isinstance(model, nn.Module):
raise RuntimeError("Class '{}' is not a PyTorch nn Module".format(
type(model).__name__))
raise RuntimeError(
"Class '{}' is not a PyTorch nn Module".format(
type(model).__name__
)
)
self.model = model
self.setter = setter
@@ -511,18 +543,21 @@ class Trainer(object):
Name of file to output tensorboard_logger to
"""
def __init__(self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None):
def __init__(
self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None
):
self.model, self.model_fn, self.optimizer, self.lr_scheduler, self.bnm_scheduler = (
model, model_fn, optimizer, lr_scheduler, bnm_scheduler)
model, model_fn, optimizer, lr_scheduler, bnm_scheduler
)
self.checkpoint_name, self.best_name = checkpoint_name, best_name
self.eval_frequency = eval_frequency
@@ -536,7 +571,8 @@ class Trainer(object):
@staticmethod
def _print(mode, epoch, loss, eval_dict, count):
to_print = "[{:d}] {}\tMean Loss: {:.4e}".format(
epoch, mode, loss / count)
epoch, mode, loss / count
)
for k, v in natsorted(eval_dict.items(), key=itemgetter(0)):
to_print += "\tMean {}: {:2.3f}%".format(k, stats.mean(v) * 1e2)
@@ -574,7 +610,8 @@ class Trainer(object):
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Training {}".format(k), 1.0 - v, step=idx)
"Training {}".format(k), 1.0 - v, step=idx
)
d_loader.dataset.randomize()
@@ -593,7 +630,8 @@ class Trainer(object):
self.optimizer.zero_grad()
_, loss, eval_res = self.model_fn(
self.model, data, eval=True, epoch=epoch)
self.model, data, eval=True, epoch=epoch
)
total_loss += loss.data[0]
count += 1
@@ -606,8 +644,7 @@ class Trainer(object):
tb_log.log_value("Eval loss", loss.data[0], step=idx)
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Eval {}".format(k), 1.0 - v, step=idx)
tb_log.log_value("Eval {}".format(k), 1.0 - v, step=idx)
d_loader.dataset.randomize()
@@ -615,12 +652,14 @@ class Trainer(object):
return total_loss / count, eval_dict
def train(self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0):
def train(
self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0
):
r"""
Call to begin training the model
@@ -649,10 +688,12 @@ class Trainer(object):
is_best = val_loss < best_loss
best_loss = min(best_loss, val_loss)
save_checkpoint(
checkpoint_state(self.model, self.optimizer, val_loss,
epoch),
checkpoint_state(
self.model, self.optimizer, val_loss, epoch
),
is_best,
filename=self.checkpoint_name,
bestname=self.best_name)
bestname=self.best_name
)
return best_loss