From 803d7e1fc61536c846c811bdee158bd21db36779 Mon Sep 17 00:00:00 2001 From: erikwijmans Date: Tue, 26 Dec 2017 19:49:52 -0500 Subject: [PATCH] Updates and some refactoring --- models/Pointnet2Cls.py | 169 ++++++++++++++++++++++--------------- models/Pointnet2SemSeg.py | 141 +++++++++++++++++-------------- models/TransformNets.py | 75 ---------------- models/__init__.py | 6 +- train_cls.py | 33 ++++---- train_sem_seg.py | 2 +- utils/CMakeLists.txt | 4 +- utils/build_ffi.py | 3 +- utils/data_utils.py | 2 +- utils/pointnet2_modules.py | 2 +- 10 files changed, 211 insertions(+), 226 deletions(-) delete mode 100644 models/TransformNets.py diff --git a/models/Pointnet2Cls.py b/models/Pointnet2Cls.py index 3d1f3a8..b598626 100644 --- a/models/Pointnet2Cls.py +++ b/models/Pointnet2Cls.py @@ -1,96 +1,133 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import Variable - import os, sys BASE_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(BASE_DIR) -sys.path.append(os.path.join(BASE_DIR, "..", "utils")) +sys.path.append(os.path.join(BASE_DIR, "../utils")) +import torch +import torch.nn as nn +from torch.autograd import Variable import pytorch_utils as pt_utils -from TransformNets import TransformNet, TranslationNet +from pointnet2_modules import PointnetSAModule, PointnetFPModule, PointnetSAModuleMSG +from pointnet2_utils import RandomDropout +from collections import namedtuple def model_fn_decorator(criterion): - transform_reg = 1e-3 + ModelReturn = namedtuple("ModelReturn", ['preds', 'loss', 'acc']) - def ortho_loss(matrix): - return torch.dist( - matrix.bmm(matrix.transpose(1, 2)), - Variable( - torch.eye(matrix.size(1), matrix.size(2)).type( - torch.cuda.FloatTensor))) + def model_fn(model, data, epoch=0, eval=False): + inputs, labels = data + inputs = Variable(inputs.cuda(async=True), volatile=eval) + labels = Variable(labels.cuda(async=True), volatile=eval) - def wrapped(model, inputs, labels): - labels = labels.squeeze() - preds, end_points = model(inputs) + xyz = inputs[..., :3] + if inputs.size(2) > 3: + points = inputs[..., 3:] + else: + points = None - transform_loss = 0.0 - for _, T in end_points.items(): - transform_loss += ortho_loss(T) + preds = model(xyz, points) + labels = labels.view(-1) + loss = criterion(preds, labels) - preds_loss = criterion(preds, labels) - loss = preds_loss + transform_reg * transform_loss + _, classes = torch.max(preds.data, -1) + acc = (classes == labels.data).sum() / labels.numel() - _, classes = torch.max(preds, 1) - acc = (classes == labels).sum() + return ModelReturn(preds, loss, {"acc": acc}) - return preds, loss, acc.data[0] - - return wrapped + return model_fn -class PointnetCls(nn.Module): - def __init__(self): +class Pointnet2SSG(nn.Module): + def __init__(self, num_classes, input_channels=9): super().__init__() - self.translation_net = TranslationNet() - self.t_net = TransformNet(1, 3, 3, scale=False) - self.f_net = TransformNet(64, 1, 64, scale=False) + self.SA_modules = nn.ModuleList() + self.SA_modules.append( + PointnetSAModule( + npoint=512, + radius=0.2, + nsample=64, + mlp=[input_channels, 64, 64, 128])) + self.SA_modules.append( + PointnetSAModule( + npoint=128, + radius=0.4, + nsample=64, + mlp=[128 + 3, 128, 128, 256])) + self.SA_modules.append(PointnetSAModule(mlp=[256 + 3, 256, 512, 1024])) - self.input_mlp = nn.Sequential( - pt_utils.Conv2d(1, 64, [1, 3], bn=True), - pt_utils.Conv2d(64, 64, bn=True)) - - self.second_mlp = pt_utils.SharedMLP([64, 64, 128, 1024], bn=True) - - self.final_mlp = nn.Sequential( + 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(0.3), pt_utils.FC(256, 40, activation=None)) + nn.Dropout(p=0.5), + pt_utils.FC(256, num_classes, activation=None)) - def forward(self, points: torch.Tensor): - batch_size, n_points, _ = points.size() - end_points = {} + def forward(self, xyz, points=None): + for module in self.SA_modules: + xyz, points = module(xyz, points) - points = points + self.translation_net(points).unsqueeze(1) - points, transform = self.apply_transform( - points, *self.t_net(points.unsqueeze(1))) - - points = self.input_mlp(points.unsqueeze(1)) - - points, transform = self.apply_transform(points.squeeze().transpose( - 1, 2), *self.f_net(points)) - end_points['trans2'] = transform - - points = F.max_pool2d( - self.second_mlp(points.transpose(1, 2).unsqueeze(-1)), - kernel_size=[n_points, 1]) - return self.final_mlp(points.view(-1, 1024)), end_points + return self.FC_layer(points.squeeze(1)) - def apply_transform(self, points, rotation, scale=None): - points = points @ rotation - if scale is not None: - points = points * scale.contiguous().view(-1, 1, 1).repeat( - 1, points.size(1), points.size(2)) +class Pointnet2MSG(nn.Module): + def __init__(self, num_classes, input_channels=9): + super().__init__() - return points, rotation + self.SA_modules = nn.ModuleList() + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=512, + radii=[0.1, 0.2, 0.4], + nsamples=[32, 64, 128], + mlps=[[input_channels, 32, 32, + 64], [input_channels, 64, 64, 128], + [input_channels, 64, 96, 128]])) + + input_channels = 64 + 128 + 128 + 3 + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=128, + 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]])) + self.SA_modules.append( + 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)) + + def forward(self, xyz, points=None): + for module in self.SA_modules: + xyz, points = module(xyz, points) + + return self.FC_layer(points.squeeze(1)) if __name__ == "__main__": from torch.autograd import Variable - model = PointnetCls() - data = Variable(torch.randn(2, 10, 3)) - print(model(data)) + import numpy as np + import torch.optim as optim + B = 2 + N = 32 + inputs = torch.randn(B, N, 9).cuda() + labels = torch.from_numpy(np.random.randint(0, 3, size=B)).cuda() + model = Pointnet2SSG(3) + model.cuda() + + optimizer = optim.Adam(model.parameters(), lr=1e-2) + + model_fn = model_fn_decorator(nn.CrossEntropyLoss()) + for _ in range(20): + optimizer.zero_grad() + _, loss, _ = model_fn(model, (inputs, labels)) + loss.backward() + print(loss.data[0]) + optimizer.step() diff --git a/models/Pointnet2SemSeg.py b/models/Pointnet2SemSeg.py index 90f0726..b076331 100644 --- a/models/Pointnet2SemSeg.py +++ b/models/Pointnet2SemSeg.py @@ -43,22 +43,34 @@ class Pointnet2SSG(nn.Module): self.initial_dropout = RandomDropout(0.4) - self.SA_module0 = PointnetSAModule( - npoint=1024, - radius=0.1, - nsample=32, - mlp=[input_channels, 32, 32, 64]) - self.SA_module1 = PointnetSAModule( - npoint=256, radius=0.2, nsample=32, mlp=[64 + 3, 64, 64, 128]) - self.SA_module2 = PointnetSAModule( - npoint=64, radius=0.4, nsample=32, mlp=[128 + 3, 128, 128, 256]) - self.SA_module3 = PointnetSAModule( - npoint=16, radius=0.8, nsample=32, mlp=[256 + 3, 256, 256, 512]) + self.SA_modules = nn.ModuleList() + self.SA_modules.append( + PointnetSAModule( + npoint=1024, + radius=0.1, + nsample=32, + mlp=[input_channels, 32, 32, 64])) + self.SA_modules.append( + PointnetSAModule( + 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])) + self.SA_modules.append( + PointnetSAModule( + npoint=16, + radius=0.8, + nsample=32, + mlp=[256 + 3, 256, 256, 512])) - self.FP_module0 = PointnetFPModule(mlp=[512 + 256, 256, 256]) - self.FP_module1 = PointnetFPModule(mlp=[256 + 128, 256, 256]) - self.FP_module2 = PointnetFPModule(mlp=[256 + 64, 256, 128]) - self.FP_module3 = PointnetFPModule(mlp=[128 + 6, 128, 128, 128]) + self.FP_modules = nn.ModuleList() + 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(), @@ -72,18 +84,17 @@ class Pointnet2SSG(nn.Module): l0_xyz = self.initial_dropout(xyz) l0_points = None - l1_xyz, l1_points = self.SA_module0(l0_xyz, l0_points) - l2_xyz, l2_points = self.SA_module1(l1_xyz, l1_points) - l3_xyz, l3_points = self.SA_module2(l2_xyz, l2_points) - l4_xyz, l4_points = self.SA_module3(l3_xyz, l3_points) + l_xyz, l_points = [l0_xyz], [l0_points] + for i in range(len(self.SA_modules)): + li_xyz, li_points = self.SA_modules[i](l_xyz[i], l_points[i]) + l_xyz.append(li_xyz) + l_points.append(li_points) - l3_points = self.FP_module0(l3_xyz, l4_xyz, l3_points, l4_points) - l2_points = self.FP_module1(l2_xyz, l3_xyz, l2_points, l3_points) - l1_points = self.FP_module2(l1_xyz, l2_xyz, l1_points, l2_points) - l0_points = self.FP_module3(l0_xyz, l1_xyz, l0_points, - l1_points).transpose(1, 2) + 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]) - return self.FC_layer(l0_points).transpose(1, 2).contiguous() + return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous() class Pointnet2MSG(nn.Module): @@ -93,43 +104,50 @@ class Pointnet2MSG(nn.Module): self.initial_dropout = RandomDropout(0.95, inplace=True) self.initial_dropout = None + self.SA_modules = nn.ModuleList() c_in = input_channels - self.SA_module0 = PointnetSAModuleMSG( - npoint=1024, - radii=[0.05, 0.1], - nsamples=[16, 32], - mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]]) + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=1024, + radii=[0.05, 0.1], + nsamples=[16, 32], + mlps=[[c_in, 16, 16, 32], [c_in, 32, 32, 64]])) c_out_0 = 32 + 64 c_in = c_out_0 + 3 - self.SA_module1 = PointnetSAModuleMSG( - npoint=256, - radii=[0.1, 0.2], - nsamples=[16, 32], - mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]]) + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=256, + radii=[0.1, 0.2], + nsamples=[16, 32], + mlps=[[c_in, 64, 64, 128], [c_in, 64, 96, 128]])) c_out_1 = 128 + 128 c_in = c_out_1 + 3 - self.SA_module2 = PointnetSAModuleMSG( - npoint=64, - radii=[0.2, 0.4], - nsamples=[16, 32], - mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]]) + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=64, + radii=[0.2, 0.4], + nsamples=[16, 32], + mlps=[[c_in, 128, 196, 256], [c_in, 128, 196, 256]])) c_out_2 = 256 + 256 c_in = c_out_2 + 3 - self.SA_module3 = PointnetSAModuleMSG( - npoint=16, - radii=[0.4, 0.8], - nsamples=[16, 32], - mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]]) + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=16, + radii=[0.4, 0.8], + nsamples=[16, 32], + mlps=[[c_in, 256, 256, 512], [c_in, 256, 384, 512]])) c_out_3 = 512 + 512 - self.FP_module3 = PointnetFPModule(mlp=[c_out_3 + c_out_2, 512, 512]) - self.FP_module2 = PointnetFPModule(mlp=[512 + c_out_1, 512, 512]) - self.FP_module1 = PointnetFPModule(mlp=[512 + c_out_0, 256, 256]) - self.FP_module0 = PointnetFPModule( - mlp=[256 + input_channels - 3, 128, 128]) + self.FP_modules = nn.ModuleList() + self.FP_modules.append( + 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])) self.FC_layer = nn.Sequential( pt_utils.Conv1d(128, 128, bn=True), nn.Dropout(), @@ -142,20 +160,17 @@ class Pointnet2MSG(nn.Module): elif self.initial_dropout is not None: xyz = self.initial_dropout(xyz) - l0_xyz, l0_points = xyz, points + l_xyz, l_points = [xyz], [points] + for i in range(len(self.SA_modules)): + li_xyz, li_points = self.SA_modules[i](l_xyz[i], l_points[i]) + l_xyz.append(li_xyz) + l_points.append(li_points) - l1_xyz, l1_points = self.SA_module0(l0_xyz, l0_points) - l2_xyz, l2_points = self.SA_module1(l1_xyz, l1_points) - l3_xyz, l3_points = self.SA_module2(l2_xyz, l2_points) - l4_xyz, l4_points = self.SA_module3(l3_xyz, l3_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]) - l3_points = self.FP_module3(l3_xyz, l4_xyz, l3_points, l4_points) - l2_points = self.FP_module2(l2_xyz, l3_xyz, l2_points, l3_points) - l1_points = self.FP_module1(l1_xyz, l2_xyz, l1_points, l2_points) - l0_points = self.FP_module0(l0_xyz, l1_xyz, l0_points, - l1_points).transpose(1, 2) - - return self.FC_layer(l0_points).transpose(1, 2).contiguous() + return self.FC_layer(l_points[0].transpose(1, 2)).transpose(1, 2).contiguous() if __name__ == "__main__": @@ -170,7 +185,7 @@ if __name__ == "__main__": model = Pointnet2MSG(3) model.cuda() - optimizer = optim.Adam(model.parameters(), lr=1e-5) + optimizer = optim.Adam(model.parameters(), lr=1e-2) model_fn = model_fn_decorator(nn.CrossEntropyLoss()) for _ in range(20): diff --git a/models/TransformNets.py b/models/TransformNets.py deleted file mode 100644 index 14220d9..0000000 --- a/models/TransformNets.py +++ /dev/null @@ -1,75 +0,0 @@ -import torch -import torch.nn as nn -from torch.autograd import Variable -import torch.nn.functional as F - -import os, sys -BASE_DIR = os.path.dirname(os.path.abspath(__file__)) -sys.path.append(BASE_DIR) - -import pytorch_utils as pt_utils - - -class TransformNet(nn.Module): - def __init__(self, in_size, channels, K, scale=False): - super().__init__() - self.K, self.scale = K, scale - - self.convs = nn.Sequential() - self.convs.add_module('conv0', - pt_utils.Conv2d( - in_size, 64, kernel_size=[1, channels], bn=True)) - self.convs.add_module('rest', - pt_utils.SharedMLP([64, 128, 1024], bn=True)) - - self.fc = nn.Sequential( - pt_utils.FC(1024, 512, bn=True), pt_utils.FC(512, 256, bn=True)) - - outsize = K * K - if scale: - outsize += 1 - - self.final_W = nn.Parameter(torch.FloatTensor(256, outsize)) - self.final_b = nn.Parameter(torch.FloatTensor(outsize)) - - self.init_weights() - - def forward(self, X): - X = self.convs(X) - X = F.adaptive_max_pool2d(X, [1, 1]) - X = self.fc(X.view(-1, 1024)) - X = X @ self.final_W + self.final_b - - rotation = X[:, 0:self.K * self.K].contiguous().view( - -1, self.K, self.K) - - if not self.scale: - return rotation, None - - scale = X[:, -1].contiguous() - - return rotation, scale - - def init_weights(self): - torch.nn.init.constant(self.final_W, 0) - self.final_b.data[:self.K * self.K] = (torch.eye( - self.K, self.K) + 1e-1 * torch.randn(self.K, self.K)).view(-1) - if self.scale: - self.final_b.data[-1] = 1.0 - - -class TranslationNet(nn.Module): - def forward(self, X): - return -torch.mean(X, dim=1) - - -if __name__ == "__main__": - from torch.autograd import Variable - net = TransformNet(5, 1, 3, True) - net.init_weights() - data = Variable(torch.FloatTensor(1, 5, 10, 1)) - print(net(data)) - - net = TranslationNet(5, 1, 3) - net.init_weights() - print(net(data)) diff --git a/models/__init__.py b/models/__init__.py index 73dd71a..3173aa0 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1 +1,5 @@ -from .Pointnet2SemSeg import Pointnet2MSG, Pointnet2SSG +from .Pointnet2SemSeg import Pointnet2MSG as Pointnet2SemMSG +from .Pointnet2SemSeg import Pointnet2SSG as Pointnet2SemSSG + +from .Pointnet2Cls import Pointnet2MSG as Pointnet2ClsMSG +from .Pointnet2Cls import Pointnet2SSG as Pointnet2ClsSSG diff --git a/train_cls.py b/train_cls.py index 09608c4..887adb5 100644 --- a/train_cls.py +++ b/train_cls.py @@ -9,49 +9,51 @@ from torchvision import transforms import os import tensorboard_logger as tb_log -from models import PointnetCls as Pointnet -from models.PointnetCls import model_fn_decorator +from models import Pointnet2ClsMSG as Pointnet +from models.Pointnet2Cls import model_fn_decorator from data import ModelNet40Cls import utils.pytorch_utils as pt_utils import utils.data_utils as d_utils import argparse def parse_args(): - parser = argparse.ArgumentParser(description="Arg parser") + parser = argparse.ArgumentParser( + description="Arguments for cls training", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( - "-batch_size", type=int, default=128, help="Batch size [default: 128]") + "-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 [default: 1024]") + 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", type=float, default=1e-2, - help="Initial learning rate [default: 1e-2]") + help="Initial learning rate") parser.add_argument( "-lr_decay", type=float, default=0.7, - help="Learning rate decay gamma [default: 0.7]") + help="Learning rate decay gamma") parser.add_argument( "-decay_step", type=int, default=20, - help="Learning rate decay step [default: 20]") + help="Learning rate decay step") parser.add_argument( "-bn_momentum", type=float, default=0.5, - help="Initial batch norm momentum [default: 0.5]") + help="Initial batch norm momentum") parser.add_argument( "-bnm_decay", type=float, default=0.5, - help="Batch norm momentum decay gamma [default: 0.5]") + help="Batch norm momentum decay gamma") parser.add_argument( "-checkpoint", type=str, default=None, help="Checkpoint to start from") parser.add_argument( @@ -74,8 +76,7 @@ if __name__ == "__main__": transforms = transforms.Compose([ d_utils.PointcloudToTensor(), - d_utils.PointcloudRotate(x_axis=True), - d_utils.PointcloudScale(), + d_utils.PointcloudRotate(x_axis=True, z_axis=True), d_utils.PointcloudTranslate(), d_utils.PointcloudJitter() ]) @@ -99,7 +100,7 @@ if __name__ == "__main__": tb_log.configure('runs/{}'.format(args.run_name)) - model = Pointnet() + model = Pointnet(input_channels=3, num_classes=40) model.cuda() optimizer = optim.Adam( model.parameters(), lr=args.lr, weight_decay=args.weight_decay) @@ -107,7 +108,7 @@ if __name__ == "__main__": 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_prec = pt_utils.load_checkpoint( + start_epoch, best_loss = pt_utils.load_checkpoint( model, optimizer, filename=args.checkpoint.split(".")[0]) lr_scheduler = lr_sched.LambdaLR( @@ -118,7 +119,7 @@ if __name__ == "__main__": lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lambda=lr_lbmd) bnm_scheduler = pt_utils.BNMomentumScheduler(model, bn_lambda=bn_lbmd) - best_prec = 0.0 + best_loss = 1e10 start_epoch = 1 model_fn = model_fn_decorator(nn.CrossEntropyLoss()) @@ -137,7 +138,7 @@ if __name__ == "__main__": args.epochs, train_loader, test_loader, - best_prec=best_prec) + best_loss=best_loss) if start_epoch == args.epochs: _ = trainer.eval_epoch(start_epoch, test_loader) diff --git a/train_sem_seg.py b/train_sem_seg.py index 585c9c5..76d1cb0 100644 --- a/train_sem_seg.py +++ b/train_sem_seg.py @@ -9,7 +9,7 @@ import numpy as np import tensorboard_logger as tb_log import os -from models import Pointnet2MSG as Pointnet +from models import Pointnet2SSG as Pointnet from models.Pointnet2SemSeg import model_fn_decorator from data import Indoor3DSemSeg import utils.pytorch_utils as pt_utils diff --git a/utils/CMakeLists.txt b/utils/CMakeLists.txt index ae60ded..067d46e 100644 --- a/utils/CMakeLists.txt +++ b/utils/CMakeLists.txt @@ -1,7 +1,7 @@ project(PointNet2) cmake_minimum_required(VERSION 3.5) -find_package(CUDA) +find_package(CUDA REQUIRED) include_directories("${CMAKE_SOURCE_DIR}/cinclude") cuda_include_directories("${CMAKE_SOURCE_DIR}/cinclude") @@ -9,11 +9,13 @@ file(GLOB cuda_kernels_src "csrc/*.cu") cuda_compile(cuda_kernels SHARED ${cuda_kernels_src} OPTIONS -O3) file(GLOB wrapper_headers "cinclude/*wrapper.h") +file(GLOB wrapper_sources "csrs/*.c") add_custom_command(OUTPUT "${CMAKE_SOURCE_DIR}/_ext/__ext.so" WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} COMMAND python "${CMAKE_SOURCE_DIR}/build_ffi.py" ${cuda_kernels} DEPENDS ${cuda_kernels} DEPENDS ${wrapper_headers} + DEPENDS ${wrapper_sources} VERBATIM) add_custom_target(ext ALL diff --git a/utils/build_ffi.py b/utils/build_ffi.py index 83dc21a..2983d81 100644 --- a/utils/build_ffi.py +++ b/utils/build_ffi.py @@ -16,7 +16,8 @@ ffi = create_extension( relative_to=__file__, with_cuda=True, extra_objects=extra_objects, - include_dirs=[path.join(base_dir, 'cinclude')]) + include_dirs=[path.join(base_dir, 'cinclude')], + verbose=False) if __name__ == "__main__": assert torch.cuda.is_available(), "Needs CUDA!" diff --git a/utils/data_utils.py b/utils/data_utils.py index 62f23e0..168025e 100644 --- a/utils/data_utils.py +++ b/utils/data_utils.py @@ -16,7 +16,7 @@ class PointcloudScale(object): class PointcloudRotate(object): def __init__(self, x_axis=False, z_axis=True): assert x_axis or z_axis - self.x, self.y = x_axis, z_axis + self.x, self.z = x_axis, z_axis def _get_angles(self): rotation_angle = np.random.uniform() * 2 * np.pi diff --git a/utils/pointnet2_modules.py b/utils/pointnet2_modules.py index 6a91b25..7900f0b 100644 --- a/utils/pointnet2_modules.py +++ b/utils/pointnet2_modules.py @@ -142,7 +142,7 @@ class PointnetSAModule(nn.Module): new_xyz = pointnet2_utils.gather_points( xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint)) else: - new_xyz = xyz.new([[[0, 0, 0]]]).expand(xyz.size(0), 1, 3) + 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)