Further developments with MNIST testing.

mnist/model.py

@@ -2,14 +2,15 @@ import math
 import data_utils
 import time
 
+import numpy as np
 import torch
 from torch import nn
 from torch.autograd import Variable
 from torch.utils.data import Dataset, DataLoader
 
 from pointcnn.core import rPointCNN
-from pointcnn.util import knn_indices_func
-from pointcnn.layers import MLP
+from pointcnn.util import knn_indices_func_gpu
+from pointcnn.layers import Dense
 
 x = 2
 
@@ -36,7 +37,7 @@ class mnist_dataset(Dataset):
 
 # C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
 # (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
-paPointCNN = lambda a,b,c,d,e: rPointCNN(a, b, 3, c, d, e, knn_indices_func)
+paPointCNN = lambda a,b,c,d,e: rPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
 
 class Classifier(nn.Module):
     
@@ -52,13 +53,9 @@ class Classifier(nn.Module):
         )
 
         self.fcn = nn.Sequential(
-            nn.Linear(160, 128),
-            nn.ReLU(),
-            nn.Dropout(0.0),
-            nn.Linear(128,  64),  # throw in some batch normalization
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(64, 10),  # 10 digits
+            Dense(160, 128),
+            Dense(128,  64, drop_rate = 0.5),
+            Dense(64, 10, activation = None)
         )
 
         self.log_softmax = nn.LogSoftmax()
@@ -66,9 +63,10 @@ class Classifier(nn.Module):
     def forward(self, x):
         x = self.pcnn(x)[1]  # grab features
         logits = self.fcn(x)
-        logits = torch.mean(logits, dim = 1)
-        log_probs = self.log_softmax(logits)
-        return log_probs
+        # logits = torch.mean(logits, dim = 1)
+        return logits
+        # log_probs = self.log_softmax(logits)
+        # return log_probs
 
 model = Classifier().cuda()
 
@@ -94,14 +92,22 @@ point_num = data_train.shape[1]
 batch_num_per_epoch = int(math.ceil(num_train / batch_size))
 batch_num = batch_num_per_epoch * num_epochs
 
-dataset = mnist_dataset(data_train, label_train)
-loader = DataLoader(dataset, batch_size = batch_size)
+training_set = mnist_dataset(data_train, label_train)
+training_loader = DataLoader(training_set, batch_size = batch_size)
 
-optimizer = torch.optim.SGD(model.parameters(), lr = 0.1, momentum = 0.9)
+testing_set = mnist_dataset(data_val, label_val)
+testing_loader = DataLoader(testing_set, batch_size = 1)
+
+optimizer = torch.optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)
 loss_fn = nn.NLLLoss()
 
 for _ in range(num_epochs):
-    for data, label in loader:
+
+    n = 0
+
+    for data, label in training_loader:
+
+        n += 1
 
         data = Variable(data).cuda()
         label = Variable(label.long()).cuda()
@@ -113,8 +119,34 @@ for _ in range(num_epochs):
         t0 = time.time()
         out = model((P, F))
 
+        print(out)
+
         loss = loss_fn(out, label)
         loss.backward()
         optimizer.step()
         
-        # print(loss.data[0])
+        print("loss:", loss.data[0])
+
+        if n % 25 == 0:
+            # Testing accuracy
+            num_testing = 0
+            total = 0
+            correct = 0
+            for data, label in testing_loader:
+                if num_testing > 100:
+                    break
+                else:
+                    num_testing += 1
+                data = Variable(data).cuda()
+                label = Variable(label.long()).cuda()
+                P = data[:,:,:3]
+                F = data[:,:,3:]
+                out = model((P, F))
+                probs = nn.Softmax()(out)
+                # print(probs)
+                _, pred = probs.max(1)
+                total += 1
+                if pred.cpu().data[0] == label.cpu().data[0]:
+                    correct += 1
+            accuracy = correct / total
+            print("accuracy:", accuracy)

pointcnn/core.py

@@ -7,11 +7,11 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 try:
-    from .util import knn_indices_func
+    from .util import knn_indices_func, knn_indices_func_gpu
     from .layers import MLP, LayerNorm, Conv, SepConv, endchannels
     # from .context import timed
 except SystemError:
-    from util import knn_indices_func
+    from util import knn_indices_func, knn_indices_func_gpu
     from layers import MLP, LayerNorm, Conv, SepConv, endchannels
     # from context import timed
 
@@ -120,7 +120,6 @@ class XConv(nn.Module):
         # Weight and permute F_cat with the learned X.
         F_X = torch.matmul(X, F_cat)
         F_p = self.end_conv(F_X).squeeze(dim = 2)
-        time.sleep(5)
         return F_p
 
 class PointCNN(nn.Module):
@@ -194,11 +193,11 @@ class PointCNN(nn.Module):
         :return:
         """
         ps, P, F = x
-        P_idx = self.r_indices_func(ps.cpu(), P.cpu(), self.x_conv.N_neighbors, self.dilution).cuda()  # This step takes ~97% of the time.
+        P_idx = self.r_indices_func(ps, P, self.x_conv.N_neighbors, self.dilution)  # This step takes ~97% of the time.
         P_regional = self.select_region(P, P_idx)  # Prime target for optimization: KNN on GPU.
         if False:
             # Draw neighborhood points, for debugging.
-            t = 15
+            t = 10
             n = 0
             test_point = ps[n,t,:].cpu().data.numpy()
             neighborhood = P_regional[n,t,:,:].cpu().data.numpy()
@@ -233,45 +232,27 @@ class rPointCNN(nn.Module):
 if __name__ == "__main__":
     np.random.seed(0)
 
-    TESTING = PointCNN
+    N = 1
+    num_points = 500
+    N_rep = 20
+    D = 2
+    C_in = 16
+    C_out = 32
+    N_neighbors = 30
+    dilution = 1
 
-    if TESTING == XConv:
-        N = 4
-        N_rep = 150
-        D = 3
-        C_in = 8
-        C_out = 32
-        N_neighbors = 10
+    model = PointCNN(C_in, C_out, D, N_neighbors, dilution, N_rep, knn_indices_func_gpu).cuda()
 
-        model = XConv(C_in, C_out, D, N_neighbors, N_rep)
-        p = Variable(torch.from_numpy(np.random.rand(N,N_rep,D).astype(np.float32)))
-        P = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,D).astype(np.float32)))
-        F = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,C_in).astype(np.float32)))
+    test_P  = np.random.rand(N,num_points,D).astype(np.float32)
+    test_F  = np.random.rand(N,num_points,C_in).astype(np.float32)
+    idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
+    test_ps = test_P[:,idx,:]
 
-        out = model(p, P, F)
+    test_P = Variable(torch.from_numpy(test_P)).cuda()
+    test_F = Variable(torch.from_numpy(test_F)).cuda()
+    test_ps = Variable(torch.from_numpy(test_ps)).cuda()
 
-    elif TESTING == PointCNN:
-        N = 4
-        num_points = 10000
-        N_rep = 5000
-        D = 3
-        C_in = 128
-        C_out = 256
-        N_neighbors = 10
-        dilution = 2
-
-        model = PointCNN(C_in, C_out, D, N_neighbors, dilution, N_rep, knn_indices_func).cuda()
-
-        test_P  = np.random.rand(N,num_points,D).astype(np.float32)
-        test_F  = np.random.rand(N,num_points,C_in).astype(np.float32)
-        idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
-        test_ps = test_P[:,idx,:]
-
-        test_P = Variable(torch.from_numpy(test_P)).cuda()
-        test_F = Variable(torch.from_numpy(test_F)).cuda()
-        test_ps = Variable(torch.from_numpy(test_ps)).cuda()
-
-        print(test_F.size())
-        for _ in range(50):
-            out = model((test_ps, test_P, test_F))
-        print(out.size())
+    print(test_F.size())
+    for _ in range(1):
+        out = model((test_ps, test_P, test_F))
+    print(out.size())

pointcnn/layers.py

@@ -30,8 +30,7 @@ class Dense(nn.Module):
         self.linear = nn.Linear(in_features, out_features)
         self.activation = activation
         # self.bn = LayerNorm(out_channels) if with_bn else None
-        if drop_rate > 0:
-            self.drop = nn.Dropout(drop_rate)
+        self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None
 
     def forward(self, x):
         x = self.linear(x)

pointcnn/util.py

@@ -1,3 +1,5 @@
+import sys, os
+
 import torch
 
 import numpy as np
@@ -5,6 +7,12 @@ from sklearn.neighbors import NearestNeighbors
 
 torch.CUDA_LAUNCH_BLOCKING = 1
 
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) 
+
+sys.path.append(os.path.join(CURRENT_DIR, "..", "lib"))
+
+import pytorch_knn_cuda
+
 try:
     # from .context import pytorch_knn_cuda
     pass
@@ -59,12 +67,28 @@ def knn_indices_func(ps, P, k, d):
     ], axis = 0)
     return torch.from_numpy(region_idx)
 
-def knn_indices_func_gpu(ps, P, k):
+# def knn_indices_func_gpu(ps, P, k, d):
+# 
+#     def single_batch_knn(p, P_particular):
+#         nbrs_f = pytorch_knn_cuda.KNearestNeighbor(d*k + 1)
+#         indices = nbrs_f(P_particular, p)[0]
+#         return indices[:,1::d]
+# 
+#     region_idx = torch.stack([
+#         single_batch_knn(p, P[n]) for n, p in enumerate(ps)
+#     ], dim = 0)
+#     return region_idx
 
-    def single_batch_knn(p, P_particular):
-        nbrs_f = pytorch_knn_cuda.KNearestNeighbor(k + 1)
-        indices = nbrs_f(P_particular, p)[0]
-        return indices[:,1:]
+def knn_indices_func_gpu(ps, P, k, d):
+
+    def single_batch_knn(qry, ref):
+        n, d = ref.size()
+        m, d = qry.size()
+        mref = ref.expand(m, n, d)
+        mqry = qry.expand(n, m, d).transpose(0, 1)
+        dist2 = torch.sum((mqry - mref)**2, 2).squeeze()
+        _, inds = torch.topk(dist2, k*d + 1, dim = 1, largest = False)
+        return inds[:,1::d]
 
     region_idx = torch.stack([
         single_batch_knn(p, P[n]) for n, p in enumerate(ps)
@@ -82,5 +106,6 @@ if __name__ == "__main__":
     test_ps = test_P[:,idx,:]
     test_P = Variable(torch.from_numpy(test_P)).cuda()
     test_ps = Variable(torch.from_numpy(test_ps)).cuda()
-    out = knn_indices_func_gpu(test_ps, test_P, 5)
+    out = knn_indices_func_gpu(test_ps, test_P, 5, 3)
+
     print(out)