Updates and some refactoring

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()

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

models/TransformNets.py

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

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

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)

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

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

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

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

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)