diff --git a/mnist/context.py b/mnist/context.py
new file mode 100644
index 0000000..126a413
--- /dev/null
+++ b/mnist/context.py
@@ -0,0 +1,4 @@
+import sys, os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from lib import pointcnn
diff --git a/mnist/data_utils.py b/mnist/data_utils.py
new file mode 100644
index 0000000..1854ef0
--- /dev/null
+++ b/mnist/data_utils.py
@@ -0,0 +1,145 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import h5py
+import plyfile
+import numpy as np
+from matplotlib import cm
+import scipy.spatial.distance as distance
+
+
+def save_ply(points, filename, colors=None, normals=None):
+    vertex = np.array([tuple(p) for p in points], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')])
+    n = len(vertex)
+    desc = vertex.dtype.descr
+
+    if normals is not None:
+        vertex_normal = np.array([tuple(n) for n in normals], dtype=[('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4')])
+        assert len(vertex_normal) == n
+        desc = desc + vertex_normal.dtype.descr
+
+    if colors is not None:
+        vertex_color = np.array([tuple(c * 255) for c in colors],
+                                dtype=[('red', 'u1'), ('green', 'u1'), ('blue', 'u1')])
+        assert len(vertex_color) == n
+        desc = desc + vertex_color.dtype.descr
+
+    vertex_all = np.empty(n, dtype=desc)
+
+    for prop in vertex.dtype.names:
+        vertex_all[prop] = vertex[prop]
+
+    if normals is not None:
+        for prop in vertex_normal.dtype.names:
+            vertex_all[prop] = vertex_normal[prop]
+
+    if colors is not None:
+        for prop in vertex_color.dtype.names:
+            vertex_all[prop] = vertex_color[prop]
+
+    ply = plyfile.PlyData([plyfile.PlyElement.describe(vertex_all, 'vertex')], text=False)
+    if not os.path.exists(os.path.dirname(filename)):
+        os.makedirs(os.path.dirname(filename))
+    ply.write(filename)
+
+
+def save_ply_property(points, property, property_max, filename, cmap_name='Set1'):
+    point_num = points.shape[0]
+    colors = np.full(points.shape, 0.5)
+    cmap = cm.get_cmap(cmap_name)
+    for point_idx in range(point_num):
+        colors[point_idx] = cmap(property[point_idx] / property_max)[:3]
+    save_ply(points, filename, colors)
+
+
+def save_ply_batch(points_batch, file_path, points_num=None):
+    batch_size = points_batch.shape[0]
+    if type(file_path) != list:
+        basename = os.path.splitext(file_path)[0]
+        ext = '.ply'
+    for batch_idx in range(batch_size):
+        point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx]
+        if type(file_path) == list:
+            save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx])
+        else:
+            save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext))
+
+
+def save_ply_property_batch(points_batch, property_batch, file_path, points_num=None, property_max=None,
+                            cmap_name='Set1'):
+    batch_size = points_batch.shape[0]
+    if type(file_path) != list:
+        basename = os.path.splitext(file_path)[0]
+        ext = '.ply'
+    property_max = np.max(property_batch) if property_max is None else property_max
+    for batch_idx in range(batch_size):
+        point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx]
+        if type(file_path) == list:
+            save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num],
+                              property_max, file_path[batch_idx], cmap_name)
+        else:
+            save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num],
+                              property_max, '%s_%04d%s' % (basename, batch_idx, ext), cmap_name)
+
+
+def save_ply_point_with_normal(data_sample, folder):
+    for idx, sample in enumerate(data_sample):
+        filename_pts = os.path.join(folder, '{:08d}.ply'.format(idx))
+        save_ply(sample[..., :3], filename_pts, normals=sample[..., 3:])
+
+
+def grouped_shuffle(inputs):
+    for idx in range(len(inputs) - 1):
+        assert (len(inputs[idx]) == len(inputs[idx + 1]))
+
+    shuffle_indices = np.arange(inputs[0].shape[0])
+    np.random.shuffle(shuffle_indices)
+    outputs = []
+    for idx in range(len(inputs)):
+        outputs.append(inputs[idx][shuffle_indices, ...])
+    return outputs
+
+
+def load_cls(filelist):
+    points = []
+    labels = []
+
+    folder = os.path.dirname(filelist)
+    for line in open(filelist):
+        filename = os.path.basename(line.rstrip())
+        data = h5py.File(os.path.join(folder, filename))
+        if 'normal' in data:
+            points.append(np.concatenate([data['data'][...], data['data'][...]], axis=-1).astype(np.float32))
+        else:
+            points.append(data['data'][...].astype(np.float32))
+        labels.append(np.squeeze(data['label'][:]).astype(np.int32))
+    return (np.concatenate(points, axis=0),
+            np.concatenate(labels, axis=0))
+
+
+def load_cls_train_val(filelist, filelist_val):
+    data_train, label_train = grouped_shuffle(load_cls(filelist))
+    data_val, label_val = load_cls(filelist_val)
+    return data_train, label_train, data_val, label_val
+
+
+def load_seg(filelist):
+    points = []
+    labels = []
+    point_nums = []
+    labels_seg = []
+
+    folder = os.path.dirname(filelist)
+    for line in open(filelist):
+        filename = os.path.basename(line.rstrip())
+        data = h5py.File(os.path.join(folder, filename))
+        points.append(data['data'][...].astype(np.float32))
+        labels.append(data['label'][...].astype(np.int32))
+        point_nums.append(data['data_num'][...].astype(np.int32))
+        labels_seg.append(data['label_seg'][...].astype(np.int32))
+    return (np.concatenate(points, axis=0),
+            np.concatenate(labels, axis=0),
+            np.concatenate(point_nums, axis=0),
+            np.concatenate(labels_seg, axis=0))
diff --git a/mnist/download_datasets.py b/mnist/download_datasets.py
new file mode 100755
index 0000000..b0ff80b
--- /dev/null
+++ b/mnist/download_datasets.py
@@ -0,0 +1,140 @@
+#!/usr/bin/python3
+'''Download datasets for this project.'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+import gzip
+import html
+import shutil
+import tarfile
+import zipfile
+import requests
+import argparse
+from tqdm import tqdm
+
+
+# from https://gist.github.com/hrouault/1358474
+def query_yes_no(question, default="yes"):
+    """Ask a yes/no question via raw_input() and return their answer.
+    "question" is a string that is presented to the user.
+    "default" is the presumed answer if the user just hits <Enter>.
+        It must be "yes" (the default), "no" or None (meaning
+        an answer is required of the user).
+    The "answer" return value is one of "yes" or "no".
+    """
+    valid = {"yes": True, "y": True, "ye": True,
+             "no": False, "n": False}
+    if default == None:
+        prompt = " [y/n] "
+    elif default == "yes":
+        prompt = " [Y/n] "
+    elif default == "no":
+        prompt = " [y/N] "
+    else:
+        raise ValueError("invalid default answer: '%s'" % default)
+
+    while True:
+        sys.stdout.write(question + prompt)
+        choice = input().lower()
+        if default is not None and choice == '':
+            return valid[default]
+        elif choice in valid:
+            return valid[choice]
+        else:
+            sys.stdout.write("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")
+
+
+def download_from_url(url, dst):
+    download = True
+    if os.path.exists(dst):
+        download = query_yes_no('Seems you have downloaded %s to %s, overwrite?' % (url, dst), default='no')
+        if download:
+            os.remove(dst)
+
+    if download:
+        response = requests.get(url, stream=True)
+        total_size = int(response.headers.get('content-length', 0))
+        chunk_size = 1024 * 1024
+        bars = total_size // chunk_size
+        with open(dst, "wb") as handle:
+            for data in tqdm(response.iter_content(chunk_size=chunk_size), total=bars, desc=url.split('/')[-1],
+                             unit='M'):
+                handle.write(data)
+
+
+def download_and_unzip(url, root, dataset):
+    folder = os.path.join(root, dataset)
+    folder_zips = os.path.join(folder, 'zips')
+    if not os.path.exists(folder_zips):
+        os.makedirs(folder_zips)
+    filename_zip = os.path.join(folder_zips, url.split('/')[-1])
+
+    download_from_url(url, filename_zip)
+
+    if filename_zip.endswith('.zip'):
+        zip_ref = zipfile.ZipFile(filename_zip, 'r')
+        zip_ref.extractall(folder)
+        zip_ref.close()
+    elif filename_zip.endswith(('.tar.gz', '.tgz')):
+        tarfile.open(name=filename_zip, mode="r:gz").extractall(folder)
+    elif filename_zip.endswith('.gz'):
+        filename_no_gz = filename_zip[:-3]
+        with gzip.open(filename_zip, 'rb') as f_in, open(filename_no_gz, 'wb') as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--folder', '-f', help='Path to data folder.')
+    parser.add_argument('--dataset', '-d', help='Dataset to download.')
+    args = parser.parse_args()
+    print(args)
+
+    root = args.folder if args.folder else '../../data'
+    if args.dataset == 'tu_berlin':
+        download_and_unzip('http://cybertron.cg.tu-berlin.de/eitz/projects/classifysketch/sketches_svg.zip', root,
+                           args.dataset)
+    elif args.dataset == 'modelnet':
+        download_and_unzip('https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip', root, args.dataset)
+        folder = os.path.join(root, args.dataset)
+        folder_h5 = os.path.join(folder, 'modelnet40_ply_hdf5_2048')
+        for filename in os.listdir(folder_h5):
+            shutil.move(os.path.join(folder_h5, filename), os.path.join(folder, filename))
+        shutil.rmtree(folder_h5)
+    elif args.dataset == 'shapenet_partseg':
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/train_data.zip', root, args.dataset)
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/train_label.zip', root, args.dataset)
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/val_data.zip', root, args.dataset)
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/val_label.zip', root, args.dataset)
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/test_data.zip', root, args.dataset)
+        download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/test_label.zip', root, args.dataset)
+    elif args.dataset == 'mnist':
+        download_and_unzip('http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz', root, args.dataset)
+        download_and_unzip('http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz', root, args.dataset)
+        download_and_unzip('http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz', root, args.dataset)
+        download_and_unzip('http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz', root, args.dataset)
+    elif args.dataset == 'cifar10':
+        download_and_unzip('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz', root, args.dataset)
+    elif args.dataset == 'quick_draw':
+        url_categories = 'https://raw.githubusercontent.com/googlecreativelab/quickdraw-dataset/master/categories.txt'
+        folder = os.path.join(root, args.dataset)
+        folder_zips = os.path.join(folder, 'zips')
+        if not os.path.exists(folder_zips):
+            os.makedirs(folder_zips)
+        filename_categories = os.path.join(folder_zips, url_categories.split('/')[-1])
+        download_from_url(url_categories, filename_categories)
+
+        categories = [line.strip() for line in open(filename_categories, 'r')]
+        url_base = 'https://storage.googleapis.com/quickdraw_dataset/sketchrnn/'
+        for category in categories:
+            url = url_base + html.escape(category) + '.npz'
+            filename_category = os.path.join(folder_zips, category + '.npz')
+            download_from_url(url, filename_category)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/mnist/model.py b/mnist/model.py
new file mode 100644
index 0000000..4dd3d35
--- /dev/null
+++ b/mnist/model.py
@@ -0,0 +1,120 @@
+import math
+import data_utils
+import time
+
+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
+
+x = 2
+
+# N_neighbors, dilution, N_rep, C_out
+# 8          , 1,        all  , 16 * x
+# 8          , 2,        all  , 32 * x
+# 8          , 4,        all  , 48 * x
+# 12         , 4,        120  , 64 * x
+# 12         , 6,        120  , 80 * x
+
+# Data_dim = 3
+
+class mnist_dataset(Dataset):
+
+    def __init__(self, data, labels):
+        self.data = data
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, i):
+        return self.data[i], self.labels[i]
+
+# 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)
+
+class Classifier(nn.Module):
+    
+    def __init__(self):
+        super(Classifier, self).__init__()
+
+        self.pcnn = nn.Sequential(
+            paPointCNN(  1,  32,  8, 1, 256),
+            paPointCNN( 32,  64,  8, 2, 256),
+            paPointCNN( 64,  96,  8, 4, 256),
+            paPointCNN( 96, 128, 12, 4, 120),
+            paPointCNN(128, 160, 12, 6, 120),
+        )
+
+        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
+        )
+
+        self.log_softmax = nn.LogSoftmax()
+
+    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
+
+model = Classifier().cuda()
+
+num_class = 10
+sample_num = 160
+batch_size = 32
+num_epochs = 2048
+jitter = 0.01
+jitter_val = 0.01
+
+rotation_range = [0, math.pi / 18, 0, 'g']
+rotation_rage_val = [0, 0, 0, 'u']
+order = 'rxyz'
+
+scaling_range = [0.05, 0.05, 0.05, 'g']
+scaling_range_val = [0, 0, 0, 'u']
+
+data_train, label_train, data_val, label_val = data_utils.load_cls_train_val("./mnist/zips/train_files.txt", "./mnist/zips/test_files.txt")
+
+num_train = data_train.shape[0]
+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)
+
+optimizer = torch.optim.SGD(model.parameters(), lr = 0.1, momentum = 0.9)
+loss_fn = nn.NLLLoss()
+
+for _ in range(num_epochs):
+    for data, label in loader:
+
+        data = Variable(data).cuda()
+        label = Variable(label.long()).cuda()
+        P = data[:,:,:3]
+        F = data[:,:,3:]
+
+        optimizer.zero_grad()
+
+        t0 = time.time()
+        out = model((P, F))
+
+        loss = loss_fn(out, label)
+        loss.backward()
+        optimizer.step()
+        
+        # print(loss.data[0])
diff --git a/mnist/pointcnn b/mnist/pointcnn
new file mode 120000
index 0000000..954cf08
--- /dev/null
+++ b/mnist/pointcnn
@@ -0,0 +1 @@
+/home/austin/ISEE/PointCNN/pointcnn
\ No newline at end of file
diff --git a/mnist/prepare_mnist_data.py b/mnist/prepare_mnist_data.py
new file mode 100755
index 0000000..b295ba9
--- /dev/null
+++ b/mnist/prepare_mnist_data.py
@@ -0,0 +1,92 @@
+#!/usr/bin/python3
+'''Convert MNIST to points.'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+import h5py
+import random
+import argparse
+import numpy as np
+from mnist import MNIST
+from datetime import datetime
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import data_utils
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--folder', '-f', help='Path to data folder')
+    parser.add_argument('--point_num', '-p', help='Point number for each sample', type=int, default=256)
+    parser.add_argument('--save_ply', '-s', help='Convert .pts to .ply', action='store_true')
+    args = parser.parse_args()
+    print(args)
+
+    batch_size = 2048
+
+    folder_mnist = args.folder if args.folder else '../../data/mnist/zips'
+    folder_pts = os.path.join(os.path.dirname(folder_mnist), 'pts')
+
+    mnist_data = MNIST(folder_mnist)
+    mnist_train_test = [(mnist_data.load_training(), 'train'), (mnist_data.load_testing(), 'test')]
+
+    data = np.zeros((batch_size, args.point_num, 4))
+    label = np.zeros((batch_size), dtype=np.int32)
+    for ((images, labels), tag) in mnist_train_test:
+        idx_h5 = 0
+        filename_filelist_h5 = os.path.join(os.path.dirname(folder_mnist), '%s_files.txt' % tag)
+        point_num_total = 0
+        with open(filename_filelist_h5, 'w') as filelist_h5:
+            for idx_img, image in enumerate(images):
+                points = []
+                pixels = []
+                for idx_pixel, pixel in enumerate(image):
+                    if pixel == 0:
+                        continue
+                    x = idx_pixel // 28
+                    z = idx_pixel % 28
+                    points.append((x, random.random() * 1e-6, z))
+                    pixels.append(pixel)
+                point_num_total = point_num_total + len(points)
+                pixels_sum = sum(pixels)
+                probs = [pixel / pixels_sum for pixel in pixels]
+                indices = np.random.choice(list(range(len(points))), size=args.point_num,
+                                           replace=(len(points) < args.point_num), p=probs)
+                points_array = np.array(points)[indices]
+                pixels_array_1d = (np.array(pixels)[indices].astype(np.float32) / 255) - 0.5
+                pixels_array = np.expand_dims(pixels_array_1d, axis=-1)
+
+                points_min = np.amin(points_array, axis=0)
+                points_max = np.amax(points_array, axis=0)
+                points_center = (points_min + points_max) / 2
+                scale = np.amax(points_max - points_min) / 2
+                points_array = (points_array - points_center) * (0.8 / scale)
+
+                if args.save_ply:
+                    filename_pts = os.path.join(folder_pts, tag, '{:06d}.ply'.format(idx_img))
+                    data_utils.save_ply(points_array, filename_pts, colors=np.tile(pixels_array, (1, 3)) + 0.5)
+
+                idx_in_batch = idx_img % batch_size
+                data[idx_in_batch, ...] = np.concatenate((points_array, pixels_array), axis=-1)
+                label[idx_in_batch] = labels[idx_img]
+                if ((idx_img + 1) % batch_size == 0) or idx_img == len(images) - 1:
+                    item_num = idx_in_batch + 1
+                    filename_h5 = os.path.join(os.path.dirname(folder_mnist), '%s_%d.h5' % (tag, idx_h5))
+                    print('{}-Saving {}...'.format(datetime.now(), filename_h5))
+                    filelist_h5.write('./%s_%d.h5\n' % (tag, idx_h5))
+
+                    file = h5py.File(filename_h5, 'w')
+                    file.create_dataset('data', data=data[0:item_num, ...])
+                    file.create_dataset('label', data=label[0:item_num, ...])
+                    file.close()
+
+                    idx_h5 = idx_h5 + 1
+        print('Average point number in each sample is : %f!' % (point_num_total / len(images)))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/mnist/train.py b/mnist/train.py
new file mode 100644
index 0000000..e4cee7d
--- /dev/null
+++ b/mnist/train.py
@@ -0,0 +1,9 @@
+import h5py
+
+f = h5py.File("./mnist/zips/train_0.h5", 'r')
+
+data = f["data"]
+label = f["label"]
+
+for r in data[0]:
+    print(r)
diff --git a/pointcnn/context.py b/pointcnn/context.py
index c26bc1f..930c062 100644
--- a/pointcnn/context.py
+++ b/pointcnn/context.py
@@ -1,5 +1,5 @@
 import sys, os
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from lib import timed
-from lib import pytorch_knn_cuda
+# from lib import timed
+# from lib import pytorch_knn_cuda
diff --git a/pointcnn/core.py b/pointcnn/core.py
index b3013be..da0225a 100644
--- a/pointcnn/core.py
+++ b/pointcnn/core.py
@@ -8,19 +8,20 @@ import matplotlib.pyplot as plt
 
 try:
     from .util import knn_indices_func
-    from .layers import MLP, LayerNorm, DepthwiseSeparableConv2d, endchannels
-    from .context import timed
+    from .layers import MLP, LayerNorm, Conv, SepConv, endchannels
+    # from .context import timed
 except SystemError:
     from util import knn_indices_func
-    from layers import MLP, LayerNorm, DepthwiseSeparableConv2d, endchannels
-    from context import timed
+    from layers import MLP, LayerNorm, Conv, SepConv, endchannels
+    # from context import timed
 
 class XConv(nn.Module):
     """
     Vectorized pointwise convolution.
     """
 
-    def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 2):
+    def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None,
+                 mlp_width = 2):
         """
         :param C_in: Input dimension of the points' features.
         :param C_out: Output dimension of the representative point features.
@@ -48,47 +49,30 @@ class XConv(nn.Module):
 
         # Main dense linear layers
         self.mlp_lift = MLP([D] + [self.C_lifted] * (mlp_width - 1), batch_norm = False)
-        # self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1).cuda())
 
         # Layers to generate X
-        self.mid_conv = endchannels(nn.Sequential(
-            nn.Conv2d(D, N_neighbors**2, (1, N_neighbors)).cuda(),
-            nn.ReLU(),
-            nn.BatchNorm2d(N_neighbors**2).cuda(),
-        ))
-        self.mid_dwconv1 = endchannels(nn.Sequential(
-            DepthwiseSeparableConv2d(
-                in_channels = N_neighbors,
-                out_channels = N_neighbors**2,
-                kernel_size = (1, N_neighbors),
-                depth_multiplier = N_neighbors
-            ).cuda(),
-            nn.ReLU(),
-            nn.BatchNorm2d(N_neighbors ** 2).cuda()
-        ))
-        self.mid_dwconv2 = endchannels(nn.Sequential(
-            DepthwiseSeparableConv2d(
-                in_channels = N_neighbors,
-                out_channels = N_neighbors**2,
-                kernel_size = (1, N_neighbors),
-                depth_multiplier = N_neighbors
-            ).cuda(),
-            nn.ReLU(),
-            nn.BatchNorm2d(N_neighbors**2).cuda()
-        ))
+        self.mid_conv = endchannels(Conv(D, N_neighbors**2, (1, N_neighbors))).cuda()
+        self.mid_dwconv1 = endchannels(SepConv(
+            in_channels = N_neighbors,
+            out_channels = N_neighbors**2,
+            kernel_size = (1, N_neighbors),
+            depth_multiplier = N_neighbors
+        )).cuda()
+        self.mid_dwconv2 = endchannels(SepConv(
+            in_channels = N_neighbors,
+            out_channels = N_neighbors**2,
+            kernel_size = (1, N_neighbors),
+            depth_multiplier = N_neighbors
+        )).cuda()
 
         # Final 
         self.mlp = MLP([N_neighbors] * mlp_width, batch_norm = False)
-        self.end_conv = endchannels(nn.Sequential(
-            DepthwiseSeparableConv2d(
-                in_channels = C_lifted + C_in,
-                out_channels = C_out,
-                kernel_size = (1, N_neighbors),
-                depth_multiplier = 4
-            ).cuda(),
-            nn.ReLU(),
-            nn.BatchNorm2d(C_out).cuda()
-        ))
+        self.end_conv = endchannels(SepConv(
+            in_channels = C_lifted + C_in,
+            out_channels = C_out,
+            kernel_size = (1, N_neighbors),
+            depth_multiplier = 4
+        )).cuda()
 
     # @timed.timed
     def forward(self, x):
@@ -136,6 +120,7 @@ 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):
@@ -143,18 +128,22 @@ class PointCNN(nn.Module):
     TODO: Insert documentation
     """
 
-    def __init__(self, C_in, C_out, D, N_neighbors, N_rep, r_indices_func, C_lifted = None, mlp_width = 4):
+    def __init__(self, C_in, C_out, D, N_neighbors, dilution, N_rep,
+                 r_indices_func, C_lifted = None, mlp_width = 2):
         """
         :param C_in: Input dimension of the points' features.
         :param C_out: Output dimension of the representative point features.
         :param D: Spatial dimensionality of points.
         :param N_neighbors: Number of neighbors to convolve over.
+        :param N_rep: Number of representative points.
+        :param dilution: "Spread" of neighboring points.
         :param r_indices_func: Selector function of the type,
             INP
             ======
             ps : (N, N_rep, D) Representative points
             P  : (N, *, D) Point cloud
             N_neighbors : Number of points for each region.
+            dilution : "Spread" of neighboring points (analogous to stride).
 
             OUT
             ======
@@ -173,6 +162,7 @@ class PointCNN(nn.Module):
 
         self.r_indices_func = r_indices_func
         self.x_conv = XConv(C_in, C_out, D, N_neighbors, N_rep, C_lifted, mlp_width)
+        self.dilution = dilution
 
     def select_region(self, P, P_idx):
         """
@@ -204,12 +194,12 @@ class PointCNN(nn.Module):
         :return:
         """
         ps, P, F = x
-        P_idx = self.r_indices_func(ps.cpu(), P.cpu(), self.x_conv.N_neighbors).cuda()  # This step takes ~97% of the time.
-        P_regional = self.select_region(P, P_idx)                           # Prime target for optimization: KNN on GPU.
+        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_regional = self.select_region(P, P_idx)  # Prime target for optimization: KNN on GPU.
         if False:
             # Draw neighborhood points, for debugging.
-            t = 23
-            n = 3
+            t = 15
+            n = 0
             test_point = ps[n,t,:].cpu().data.numpy()
             neighborhood = P_regional[n,t,:,:].cpu().data.numpy()
             plt.scatter(P[n][:,0], P[n][:,1])
@@ -218,7 +208,27 @@ class PointCNN(nn.Module):
             plt.show()
         F_regional = self.select_region(F, P_idx)
         # ps, P, F_P -> ps_F
-        return self.x_conv((ps, P_regional, F_regional))
+        F_p = self.x_conv((ps, P_regional, F_regional))
+        return F_p
+
+class rPointCNN(nn.Module):
+    """ PointCNN with randomly sampled representative points. """
+
+    def __init__(self, *args, **kwargs):
+        super(rPointCNN, self).__init__()
+        self.pointcnn = PointCNN(*args, **kwargs)
+        self.N_rep = args[5]  # Exists because PointCNN requires it.
+
+    def forward(self, x):
+        P, F = x
+        if self.N_rep < P.size()[1]:
+            idx = np.random.choice(P.size()[1], self.N_rep, replace = False).tolist()
+            ps = P[:,idx,:]
+        else:
+            # All input points are representative points.
+            ps = P
+        ps_F = self.pointcnn((ps, P, F))
+        return ps, ps_F
 
 if __name__ == "__main__":
     np.random.seed(0)
@@ -242,14 +252,15 @@ if __name__ == "__main__":
 
     elif TESTING == PointCNN:
         N = 4
-        num_points = 4000
-        N_rep = 4000
+        num_points = 10000
+        N_rep = 5000
         D = 3
-        C_in = 768
-        C_out = 7
-        N_neighbors = 5
+        C_in = 128
+        C_out = 256
+        N_neighbors = 10
+        dilution = 2
 
-        model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func).cuda()
+        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)
diff --git a/pointcnn/layers.py b/pointcnn/layers.py
index 6a453bb..ef55c53 100644
--- a/pointcnn/layers.py
+++ b/pointcnn/layers.py
@@ -13,21 +13,78 @@ def endchannels(f, make_contiguous = False):
     class wrapped_layer(nn.Module):
         def __init__(self):
             super(wrapped_layer, self).__init__()
+            self.f = f
         def forward(self, x):
             x = x.permute(0,3,1,2)
-            x = f(x)
+            x = self.f(x)
             x = x.permute(0,2,3,1)
             return x
     return wrapped_layer()
 
-def DepthwiseSeparableConv2d(in_channels, out_channels, kernel_size, depth_multiplier):
+class Dense(nn.Module):
+
+    def __init__(self, in_features, out_features, drop_rate = 0, with_bn = True,
+                 activation = nn.ReLU()):
+        super(Dense, self).__init__()
+
+        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)
+
+    def forward(self, x):
+        x = self.linear(x)
+        if self.activation:
+            x = self.activation(x)
+        # if self.bn:
+        #     x = self.bn(x)
+        if self.drop:
+            x = self.drop(x)
+        return x
+
+class Conv(nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel_size, with_bn = True,
+                 activation = nn.ReLU()):
+        super(Conv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias = not with_bn)
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels) if with_bn else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
+
+class SepConv(nn.Module):
     """
     Depthwise separable convolution
     """
-    return nn.Sequential(
-        nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
-        nn.Conv2d(in_channels * depth_multiplier, out_channels, 1)
-    )
+
+    def __init__(self, in_channels, out_channels, kernel_size, depth_multiplier = 1,
+                 with_bn = True, activation = nn.ReLU()):
+        super(SepConv, self).__init__()
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
+            nn.Conv2d(in_channels * depth_multiplier, out_channels, 1, bias = not with_bn)
+        )
+
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels) if with_bn else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
 
 class LayerNorm(nn.Module):
     """
diff --git a/pointcnn/util.py b/pointcnn/util.py
index 623493c..65882e7 100644
--- a/pointcnn/util.py
+++ b/pointcnn/util.py
@@ -6,9 +6,11 @@ from sklearn.neighbors import NearestNeighbors
 torch.CUDA_LAUNCH_BLOCKING = 1
 
 try:
-    from .context import pytorch_knn_cuda
+    # from .context import pytorch_knn_cuda
+    pass
 except SystemError:
-    from context import pytorch_knn_cuda
+    # from context import pytorch_knn_cuda
+    pass
 
 def apply_along_dim(xs, f, dim):
     """
@@ -31,7 +33,7 @@ def zipwith_matmul(xs, ys):
     N = len(xs)
     return torch.stack([torch.mm(xs[i], ys[i]) for i in range(N)], dim = 0)
 
-def knn_indices_func(ps, P, k):
+def knn_indices_func(ps, P, k, d):
     """
     Indexing function based on K-Nearest Neighbors search.
     :type ps: FloatTensor (N, N_rep, D)
@@ -48,9 +50,9 @@ def knn_indices_func(ps, P, k):
     P = P.data.numpy()
 
     def single_batch_knn(p, P_particular):
-        nbrs = NearestNeighbors(k + 1, algorithm = "ball_tree").fit(P_particular)
+        nbrs = NearestNeighbors(d*k + 1, algorithm = "ball_tree").fit(P_particular)
         indices = nbrs.kneighbors(p)[1]
-        return indices[:,1:]
+        return indices[:,1::d]
 
     region_idx = np.stack([
         single_batch_knn(p, P[n]) for n, p in enumerate(ps)
diff --git a/requirements.txt b/requirements.txt
index f463940..bb532ba 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,6 @@ numpy
 libKMCUDA
 http://download.pytorch.org/whl/cu80/torch-0.3.1-cp35-cp35m-linux_x86_64.whl
 torchvision
+h5py
+mnist
+plyfile