chore: import upstream snapshot with attribution

examples/pytorch/pointcloud/bipointnet/ModelNetDataLoader.py

@@ -0,0 +1,125 @@
+import os
+import warnings
+
+import numpy as np
+from torch.utils.data import Dataset
+
+warnings.filterwarnings("ignore")
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def farthest_point_sample(point, npoint):
+    """
+    Farthest point sampler works as follows:
+    1. Initialize the sample set S with a random point
+    2. Pick point P not in S, which maximizes the distance d(P, S)
+    3. Repeat step 2 until |S| = npoint
+
+    Input:
+        xyz: pointcloud data, [N, D]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [npoint, D]
+    """
+    N, D = point.shape
+    xyz = point[:, :3]
+    centroids = np.zeros((npoint,))
+    distance = np.ones((N,)) * 1e10
+    farthest = np.random.randint(0, N)
+    for i in range(npoint):
+        centroids[i] = farthest
+        centroid = xyz[farthest, :]
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+    point = point[centroids.astype(np.int32)]
+    return point
+
+
+class ModelNetDataLoader(Dataset):
+    def __init__(
+        self,
+        root,
+        npoint=1024,
+        split="train",
+        fps=False,
+        normal_channel=True,
+        cache_size=15000,
+    ):
+        """
+        Input:
+            root: the root path to the local data files
+            npoint: number of points from each cloud
+            split: which split of the data, 'train' or 'test'
+            fps: whether to sample points with farthest point sampler
+            normal_channel: whether to use additional channel
+            cache_size: the cache size of in-memory point clouds
+        """
+        self.root = root
+        self.npoints = npoint
+        self.fps = fps
+        self.catfile = os.path.join(self.root, "modelnet40_shape_names.txt")
+
+        self.cat = [line.rstrip() for line in open(self.catfile)]
+        self.classes = dict(zip(self.cat, range(len(self.cat))))
+        self.normal_channel = normal_channel
+
+        shape_ids = {}
+        shape_ids["train"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_train.txt"))
+        ]
+        shape_ids["test"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_test.txt"))
+        ]
+
+        assert split == "train" or split == "test"
+        shape_names = ["_".join(x.split("_")[0:-1]) for x in shape_ids[split]]
+        # list of (shape_name, shape_txt_file_path) tuple
+        self.datapath = [
+            (
+                shape_names[i],
+                os.path.join(self.root, shape_names[i], shape_ids[split][i])
+                + ".txt",
+            )
+            for i in range(len(shape_ids[split]))
+        ]
+        print("The size of %s data is %d" % (split, len(self.datapath)))
+
+        self.cache_size = cache_size
+        self.cache = {}
+
+    def __len__(self):
+        return len(self.datapath)
+
+    def _get_item(self, index):
+        if index in self.cache:
+            point_set, cls = self.cache[index]
+        else:
+            fn = self.datapath[index]
+            cls = self.classes[self.datapath[index][0]]
+            cls = np.array([cls]).astype(np.int32)
+            point_set = np.loadtxt(fn[1], delimiter=",").astype(np.float32)
+            if self.fps:
+                point_set = farthest_point_sample(point_set, self.npoints)
+            else:
+                point_set = point_set[0 : self.npoints, :]
+
+            point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
+
+            if not self.normal_channel:
+                point_set = point_set[:, 0:3]
+
+            if len(self.cache) < self.cache_size:
+                self.cache[index] = (point_set, cls)
+
+        return point_set, cls

examples/pytorch/pointcloud/bipointnet/README.md

@@ -0,0 +1,43 @@
+## *BiPointNet: Binary Neural Network for Point Clouds*
+
+Created by [Haotong Qin](https://htqin.github.io/), [Zhongang Cai](https://scholar.google.com/citations?user=WrDKqIAAAAAJ&hl=en), [Mingyuan Zhang](https://scholar.google.com/citations?user=2QLD4fAAAAAJ&hl=en), Yifu Ding, Haiyu Zhao, Shuai Yi, [Xianglong Liu](http://sites.nlsde.buaa.edu.cn/~xlliu/), and [Hao Su](https://cseweb.ucsd.edu/~haosu/) from Beihang University, SenseTime, and UCSD.
+
+![prediction example](https://htqin.github.io/Imgs/ICLR/overview_v1.png)
+
+### Introduction
+
+This project is the official implementation of our accepted ICLR 2021 paper *BiPointNet: Binary Neural Network for Point Clouds* [[PDF]( https://openreview.net/forum?id=9QLRCVysdlO)]. To alleviate the resource constraint for real-time point cloud applications that run on edge devices, in this paper we present ***BiPointNet***, the first model binarization approach for efficient deep learning on point clouds. We first discover that the immense performance drop of binarized models for point clouds mainly stems from two challenges: aggregation-induced feature homogenization that leads to a degradation of information entropy, and scale distortion that hinders optimization and invalidates scale-sensitive structures. With theoretical justifications and in-depth analysis, our BiPointNet introduces Entropy-Maximizing Aggregation (EMA) to modulate the distribution before aggregation for the maximum information entropy, and Layer-wise Scale Recovery (LSR) to efficiently restore feature representation capacity. Extensive experiments show that BiPointNet outperforms existing binarization methods by convincing margins, at the level even comparable with the full precision counterpart. We highlight that our techniques are generic, guaranteeing significant improvements on various fundamental tasks and mainstream backbones, e.g., BiPointNet gives an impressive 14.7x speedup and 18.9x storage saving on real-world resource-constrained devices. Besides, our reasoning framework is dabnn.
+
+### How to Run
+
+```shell script
+python train_cls.py --model ${MODEL}
+```
+
+Here, `MODEL` has two choices: `bipointnet`  and  `bipointnet2_ssg`
+
+# Performance
+
+## Classification
+
+| Model           | Dataset    | Metric   | Score |
+| --------------- | ---------- | -------- | ----- |
+| BiPointNet      | ModelNet40 | Accuracy | 88.4  |
+| BiPointNet2_SSG | ModelNet40 | Accuracy | 83.1  |
+
+Because of the difference in implementation brought by the application of DGL, this version is even better than the original paper.
+
+### Citation
+
+If you find our work useful in your research, please consider citing:
+
+```
+@inproceedings{Qin:iclr21,
+  author    = {Haotong Qin and Zhongang Cai and Mingyuan Zhang 
+  and Yifu Ding and Haiyu Zhao and Shuai Yi 
+  and Xianglong Liu and Hao Su},
+  title     = {BiPointNet: Binary Neural Network for Point Clouds},
+  booktitle = {ICLR},
+  year      = {2021}
+}
+```

examples/pytorch/pointcloud/bipointnet/basic.py

@@ -0,0 +1,268 @@
+import dgl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.nn import Parameter
+from torch.nn.modules.utils import _single
+
+
+class BinaryQuantize(Function):
+    @staticmethod
+    def forward(ctx, input):
+        ctx.save_for_backward(input)
+        out = torch.sign(input)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input = ctx.saved_tensors
+        grad_input = grad_output
+        grad_input[input[0].gt(1)] = 0
+        grad_input[input[0].lt(-1)] = 0
+        return grad_input
+
+
+class BiLinearLSR(torch.nn.Linear):
+    def __init__(self, in_features, out_features, bias=False, binary_act=True):
+        super(BiLinearLSR, self).__init__(in_features, out_features, bias=bias)
+        self.binary_act = binary_act
+
+        # must register a nn.Parameter placeholder for model loading
+        # self.register_parameter('scale', None) doesn't register None into state_dict
+        # so it leads to unexpected key error when loading saved model
+        # hence, init scale with Parameter
+        # however, Parameter(None) actually has size [0], not [] as a scalar
+        # hence, init it using the following trick
+        self.register_parameter(
+            "scale", Parameter(torch.Tensor([0.0]).squeeze())
+        )
+
+    def reset_scale(self, input):
+        bw = self.weight
+        ba = input
+        bw = bw - bw.mean()
+        self.scale = Parameter(
+            (
+                F.linear(ba, bw).std()
+                / F.linear(torch.sign(ba), torch.sign(bw)).std()
+            )
+            .float()
+            .to(ba.device)
+        )
+        # corner case when ba is all 0.0
+        if torch.isnan(self.scale):
+            self.scale = Parameter(
+                (bw.std() / torch.sign(bw).std()).float().to(ba.device)
+            )
+
+    def forward(self, input):
+        bw = self.weight
+        ba = input
+        bw = bw - bw.mean()
+
+        if self.scale.item() == 0.0:
+            self.reset_scale(input)
+
+        bw = BinaryQuantize().apply(bw)
+        bw = bw * self.scale
+        if self.binary_act:
+            ba = BinaryQuantize().apply(ba)
+        output = F.linear(ba, bw)
+        return output
+
+
+class BiLinear(torch.nn.Linear):
+    def __init__(self, in_features, out_features, bias=True, binary_act=True):
+        super(BiLinear, self).__init__(in_features, out_features, bias=True)
+        self.binary_act = binary_act
+        self.output_ = None
+
+    def forward(self, input):
+        bw = self.weight
+        ba = input
+        bw = BinaryQuantize().apply(bw)
+        if self.binary_act:
+            ba = BinaryQuantize().apply(ba)
+        output = F.linear(ba, bw, self.bias)
+        self.output_ = output
+        return output
+
+
+class BiConv2d(torch.nn.Conv2d):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
+        super(BiConv2d, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+            padding_mode,
+        )
+
+    def forward(self, input):
+        bw = self.weight
+        ba = input
+        bw = bw - bw.mean()
+        bw = BinaryQuantize().apply(bw)
+        ba = BinaryQuantize().apply(ba)
+
+        if self.padding_mode == "circular":
+            expanded_padding = (
+                (self.padding[0] + 1) // 2,
+                self.padding[0] // 2,
+            )
+            return F.conv2d(
+                F.pad(ba, expanded_padding, mode="circular"),
+                bw,
+                self.bias,
+                self.stride,
+                _single(0),
+                self.dilation,
+                self.groups,
+            )
+        return F.conv2d(
+            ba,
+            bw,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+
+
+def square_distance(src, dst):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src**2, -1).view(B, N, 1)
+    dist += torch.sum(dst**2, -1).view(B, 1, M)
+    return dist
+
+
+def index_points(points, idx):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = (
+        torch.arange(B, dtype=torch.long)
+        .to(device)
+        .view(view_shape)
+        .repeat(repeat_shape)
+    )
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+
+class FixedRadiusNearNeighbors(nn.Module):
+    """
+    Ball Query - Find the neighbors with-in a fixed radius
+    """
+
+    def __init__(self, radius, n_neighbor):
+        super(FixedRadiusNearNeighbors, self).__init__()
+        self.radius = radius
+        self.n_neighbor = n_neighbor
+
+    def forward(self, pos, centroids):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+        """
+        device = pos.device
+        B, N, _ = pos.shape
+        center_pos = index_points(pos, centroids)
+        _, S, _ = center_pos.shape
+        group_idx = (
+            torch.arange(N, dtype=torch.long)
+            .to(device)
+            .view(1, 1, N)
+            .repeat([B, S, 1])
+        )
+        sqrdists = square_distance(center_pos, pos)
+        group_idx[sqrdists > self.radius**2] = N
+        group_idx = group_idx.sort(dim=-1)[0][:, :, : self.n_neighbor]
+        group_first = (
+            group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, self.n_neighbor])
+        )
+        mask = group_idx == N
+        group_idx[mask] = group_first[mask]
+        return group_idx
+
+
+class FixedRadiusNNGraph(nn.Module):
+    """
+    Build NN graph
+    """
+
+    def __init__(self, radius, n_neighbor):
+        super(FixedRadiusNNGraph, self).__init__()
+        self.radius = radius
+        self.n_neighbor = n_neighbor
+        self.frnn = FixedRadiusNearNeighbors(radius, n_neighbor)
+
+    def forward(self, pos, centroids, feat=None):
+        dev = pos.device
+        group_idx = self.frnn(pos, centroids)
+        B, N, _ = pos.shape
+        glist = []
+        for i in range(B):
+            center = torch.zeros((N)).to(dev)
+            center[centroids[i]] = 1
+            src = group_idx[i].contiguous().view(-1)
+            dst = centroids[i].view(-1, 1).repeat(1, self.n_neighbor).view(-1)
+
+            unified = torch.cat([src, dst])
+            uniq, inv_idx = torch.unique(unified, return_inverse=True)
+            src_idx = inv_idx[: src.shape[0]]
+            dst_idx = inv_idx[src.shape[0] :]
+
+            g = dgl.graph((src_idx, dst_idx))
+            g.ndata["pos"] = pos[i][uniq]
+            g.ndata["center"] = center[uniq]
+            if feat is not None:
+                g.ndata["feat"] = feat[i][uniq]
+            glist.append(g)
+        bg = dgl.batch(glist)
+        return bg
+
+
+class RelativePositionMessage(nn.Module):
+    """
+    Compute the input feature from neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(RelativePositionMessage, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, edges):
+        pos = edges.src["pos"] - edges.dst["pos"]
+        if "feat" in edges.src:
+            res = torch.cat([pos, edges.src["feat"]], 1)
+        else:
+            res = pos
+        return {"agg_feat": res}

examples/pytorch/pointcloud/bipointnet/bipointnet2.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from basic import (
+    BiConv2d,
+    BiLinearLSR,
+    FixedRadiusNNGraph,
+    RelativePositionMessage,
+)
+from dgl.geometry import farthest_point_sampler
+
+
+class BiPointNetConv(nn.Module):
+    """
+    Feature aggregation
+    """
+
+    def __init__(self, sizes, batch_size):
+        super(BiPointNetConv, self).__init__()
+        self.batch_size = batch_size
+        self.conv = nn.ModuleList()
+        self.bn = nn.ModuleList()
+        for i in range(1, len(sizes)):
+            self.conv.append(BiConv2d(sizes[i - 1], sizes[i], 1))
+            self.bn.append(nn.BatchNorm2d(sizes[i]))
+
+    def forward(self, nodes):
+        shape = nodes.mailbox["agg_feat"].shape
+        h = (
+            nodes.mailbox["agg_feat"]
+            .view(self.batch_size, -1, shape[1], shape[2])
+            .permute(0, 3, 2, 1)
+        )
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h, 2)[0]
+        feat_dim = h.shape[1]
+        h = h.permute(0, 2, 1).reshape(-1, feat_dim)
+        return {"new_feat": h}
+
+    def group_all(self, pos, feat):
+        """
+        Feature aggregation and pooling for the non-sampling layer
+        """
+        if feat is not None:
+            h = torch.cat([pos, feat], 2)
+        else:
+            h = pos
+        B, N, D = h.shape
+        _, _, C = pos.shape
+        new_pos = torch.zeros(B, 1, C)
+        h = h.permute(0, 2, 1).view(B, -1, N, 1)
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h[:, :, :, 0], 2)[0]  # [B,D]
+        return new_pos, h
+
+
+class BiSAModule(nn.Module):
+    """
+    The Set Abstraction Layer
+    """
+
+    def __init__(
+        self,
+        npoints,
+        batch_size,
+        radius,
+        mlp_sizes,
+        n_neighbor=64,
+        group_all=False,
+    ):
+        super(BiSAModule, self).__init__()
+        self.group_all = group_all
+        if not group_all:
+            self.npoints = npoints
+            self.frnn_graph = FixedRadiusNNGraph(radius, n_neighbor)
+        self.message = RelativePositionMessage(n_neighbor)
+        self.conv = BiPointNetConv(mlp_sizes, batch_size)
+        self.batch_size = batch_size
+
+    def forward(self, pos, feat):
+        if self.group_all:
+            return self.conv.group_all(pos, feat)
+
+        centroids = farthest_point_sampler(pos, self.npoints)
+        g = self.frnn_graph(pos, centroids, feat)
+        g.update_all(self.message, self.conv)
+
+        mask = g.ndata["center"] == 1
+        pos_dim = g.ndata["pos"].shape[-1]
+        feat_dim = g.ndata["new_feat"].shape[-1]
+        pos_res = g.ndata["pos"][mask].view(self.batch_size, -1, pos_dim)
+        feat_res = g.ndata["new_feat"][mask].view(self.batch_size, -1, feat_dim)
+        return pos_res, feat_res
+
+
+class BiPointNet2SSGCls(nn.Module):
+    def __init__(
+        self, output_classes, batch_size, input_dims=3, dropout_prob=0.4
+    ):
+        super(BiPointNet2SSGCls, self).__init__()
+        self.input_dims = input_dims
+
+        self.sa_module1 = BiSAModule(
+            512, batch_size, 0.2, [input_dims, 64, 64, 128]
+        )
+        self.sa_module2 = BiSAModule(
+            128, batch_size, 0.4, [128 + 3, 128, 128, 256]
+        )
+        self.sa_module3 = BiSAModule(
+            None, batch_size, None, [256 + 3, 256, 512, 1024], group_all=True
+        )
+
+        self.mlp1 = BiLinearLSR(1024, 512)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.drop1 = nn.Dropout(dropout_prob)
+
+        self.mlp2 = BiLinearLSR(512, 256)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.drop2 = nn.Dropout(dropout_prob)
+
+        self.mlp_out = BiLinearLSR(256, output_classes)
+
+    def forward(self, x):
+        if x.shape[-1] > 3:
+            pos = x[:, :, :3]
+            feat = x[:, :, 3:]
+        else:
+            pos = x
+            feat = None
+        pos, feat = self.sa_module1(pos, feat)
+        pos, feat = self.sa_module2(pos, feat)
+        _, h = self.sa_module3(pos, feat)
+
+        h = self.mlp1(h)
+        h = self.bn1(h)
+        h = F.relu(h)
+        h = self.drop1(h)
+        h = self.mlp2(h)
+        h = self.bn2(h)
+        h = F.relu(h)
+        h = self.drop2(h)
+
+        out = self.mlp_out(h)
+        return out

examples/pytorch/pointcloud/bipointnet/bipointnet_cls.py

@@ -0,0 +1,182 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from basic import BiLinear
+from torch.autograd import Variable
+
+offset_map = {1024: -3.2041, 2048: -3.4025, 4096: -3.5836}
+
+
+class Conv1d(nn.Module):
+    def __init__(self, inplane, outplane, Linear):
+        super().__init__()
+        self.lin = Linear(inplane, outplane)
+
+    def forward(self, x):
+        B, C, N = x.shape
+        x = x.permute(0, 2, 1).contiguous().view(-1, C)
+        x = self.lin(x).view(B, N, -1).permute(0, 2, 1).contiguous()
+        return x
+
+
+class EmaMaxPool(nn.Module):
+    def __init__(self, kernel_size, affine=True, Linear=BiLinear, use_bn=True):
+        super(EmaMaxPool, self).__init__()
+        self.kernel_size = kernel_size
+        self.bn3 = nn.BatchNorm1d(1024, affine=affine)
+        self.use_bn = use_bn
+
+    def forward(self, x):
+        batchsize, D, N = x.size()
+        if self.use_bn:
+            x = torch.max(x, 2, keepdim=True)[0] + offset_map[N]
+        else:
+            x = torch.max(x, 2, keepdim=True)[0] - 0.3
+        return x
+
+
+class BiPointNetCls(nn.Module):
+    def __init__(
+        self,
+        output_classes,
+        input_dims=3,
+        conv1_dim=64,
+        use_transform=True,
+        Linear=BiLinear,
+    ):
+        super(BiPointNetCls, self).__init__()
+        self.input_dims = input_dims
+        self.conv1 = nn.ModuleList()
+        self.conv1.append(Conv1d(input_dims, conv1_dim, Linear=Linear))
+        self.conv1.append(Conv1d(conv1_dim, conv1_dim, Linear=Linear))
+        self.conv1.append(Conv1d(conv1_dim, conv1_dim, Linear=Linear))
+
+        self.bn1 = nn.ModuleList()
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(Conv1d(conv1_dim, conv1_dim * 2, Linear=Linear))
+        self.conv2.append(Conv1d(conv1_dim * 2, conv1_dim * 16, Linear=Linear))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 2))
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 16))
+
+        self.maxpool = EmaMaxPool(conv1_dim * 16, Linear=Linear, use_bn=True)
+        self.pool_feat_len = conv1_dim * 16
+
+        self.mlp3 = nn.ModuleList()
+        self.mlp3.append(Linear(conv1_dim * 16, conv1_dim * 8))
+        self.mlp3.append(Linear(conv1_dim * 8, conv1_dim * 4))
+
+        self.bn3 = nn.ModuleList()
+        self.bn3.append(nn.BatchNorm1d(conv1_dim * 8))
+        self.bn3.append(nn.BatchNorm1d(conv1_dim * 4))
+
+        self.dropout = nn.Dropout(0.3)
+        self.mlp_out = Linear(conv1_dim * 4, output_classes)
+
+        self.use_transform = use_transform
+        if use_transform:
+            self.transform1 = TransformNet(input_dims)
+            self.trans_bn1 = nn.BatchNorm1d(input_dims)
+            self.transform2 = TransformNet(conv1_dim)
+            self.trans_bn2 = nn.BatchNorm1d(conv1_dim)
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        h = x.permute(0, 2, 1)
+        if self.use_transform:
+            trans = self.transform1(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn1(h))
+
+        for conv, bn in zip(self.conv1, self.bn1):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        if self.use_transform:
+            trans = self.transform2(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn2(h))
+
+        for conv, bn in zip(self.conv2, self.bn2):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.maxpool(h).view(-1, self.pool_feat_len)
+        for mlp, bn in zip(self.mlp3, self.bn3):
+            h = mlp(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.dropout(h)
+        out = self.mlp_out(h)
+        return out
+
+
+class TransformNet(nn.Module):
+    def __init__(self, input_dims=3, conv1_dim=64, Linear=BiLinear):
+        super(TransformNet, self).__init__()
+        self.conv = nn.ModuleList()
+        self.conv.append(Conv1d(input_dims, conv1_dim, Linear=Linear))
+        self.conv.append(Conv1d(conv1_dim, conv1_dim * 2, Linear=Linear))
+        self.conv.append(Conv1d(conv1_dim * 2, conv1_dim * 16, Linear=Linear))
+
+        self.bn = nn.ModuleList()
+        self.bn.append(nn.BatchNorm1d(conv1_dim))
+        self.bn.append(nn.BatchNorm1d(conv1_dim * 2))
+        self.bn.append(nn.BatchNorm1d(conv1_dim * 16))
+
+        # self.maxpool = nn.MaxPool1d(conv1_dim * 16)
+        self.maxpool = EmaMaxPool(conv1_dim * 16, Linear=Linear, use_bn=True)
+        self.pool_feat_len = conv1_dim * 16
+
+        self.mlp2 = nn.ModuleList()
+        self.mlp2.append(Linear(conv1_dim * 16, conv1_dim * 8))
+        self.mlp2.append(Linear(conv1_dim * 8, conv1_dim * 4))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 8))
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 4))
+
+        self.input_dims = input_dims
+        self.mlp_out = Linear(conv1_dim * 4, input_dims * input_dims)
+
+    def forward(self, h):
+        batch_size = h.shape[0]
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.maxpool(h).view(-1, self.pool_feat_len)
+        for mlp, bn in zip(self.mlp2, self.bn2):
+            h = mlp(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        out = self.mlp_out(h)
+
+        iden = Variable(
+            torch.from_numpy(
+                np.eye(self.input_dims).flatten().astype(np.float32)
+            )
+        )
+        iden = iden.view(1, self.input_dims * self.input_dims).repeat(
+            batch_size, 1
+        )
+        if out.is_cuda:
+            iden = iden.cuda()
+        out = out + iden
+        out = out.view(-1, self.input_dims, self.input_dims)
+        return out

examples/pytorch/pointcloud/bipointnet/train_cls.py

@@ -0,0 +1,175 @@
+import argparse
+import os
+import urllib
+from functools import partial
+
+import dgl
+import provider
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+from bipointnet2 import BiPointNet2SSGCls
+from bipointnet_cls import BiPointNetCls
+from dgl.data.utils import download, get_download_dir
+from ModelNetDataLoader import ModelNetDataLoader
+from torch.utils.data import DataLoader
+
+torch.backends.cudnn.enabled = False
+
+
+# from dataset import ModelNet
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", type=str, default="bipointnet")
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=200)
+parser.add_argument("--num-workers", type=int, default=0)
+parser.add_argument("--batch-size", type=int, default=32)
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+data_filename = "modelnet40_normal_resampled.zip"
+download_path = os.path.join(get_download_dir(), data_filename)
+local_path = args.dataset_path or os.path.join(
+    get_download_dir(), "modelnet40_normal_resampled"
+)
+if not os.path.exists(local_path):
+    download(
+        "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip",
+        download_path,
+        verify_ssl=False,
+    )
+    from zipfile import ZipFile
+
+    with ZipFile(download_path) as z:
+        z.extractall(path=get_download_dir())
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    loss_f = nn.CrossEntropyLoss()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label in tq:
+            data = data.data.numpy()
+            data = provider.random_point_dropout(data)
+            data[:, :, 0:3] = provider.random_scale_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.jitter_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.shift_point_cloud(data[:, :, 0:3])
+            data = torch.tensor(data)
+            label = label[:, 0]
+
+            num_examples = label.shape[0]
+            data, label = data.to(dev), label.to(dev).squeeze().long()
+            opt.zero_grad()
+            logits = net(data)
+            loss = loss_f(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            num_batches += 1
+            count += num_examples
+            loss = loss.item()
+            correct = (preds == label).sum().item()
+            total_loss += loss
+            total_correct += correct
+
+            tq.set_postfix(
+                {
+                    "AvgLoss": "%.5f" % (total_loss / num_batches),
+                    "AvgAcc": "%.5f" % (total_correct / count),
+                }
+            )
+    scheduler.step()
+
+
+def evaluate(net, test_loader, dev):
+    net.eval()
+
+    total_correct = 0
+    count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label in tq:
+                label = label[:, 0]
+                num_examples = label.shape[0]
+                data, label = data.to(dev), label.to(dev).squeeze().long()
+                logits = net(data)
+                _, preds = logits.max(1)
+
+                correct = (preds == label).sum().item()
+                total_correct += correct
+                count += num_examples
+
+                tq.set_postfix({"AvgAcc": "%.5f" % (total_correct / count)})
+
+    return total_correct / count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if args.model == "bipointnet":
+    net = BiPointNetCls(40, input_dims=6)
+elif args.model == "bipointnet2_ssg":
+    net = BiPointNet2SSGCls(40, batch_size, input_dims=6)
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+opt = optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4)
+
+scheduler = optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.7)
+
+train_dataset = ModelNetDataLoader(local_path, 1024, split="train")
+test_dataset = ModelNetDataLoader(local_path, 1024, split="test")
+train_loader = torch.utils.data.DataLoader(
+    train_dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    num_workers=num_workers,
+    drop_last=True,
+)
+test_loader = torch.utils.data.DataLoader(
+    test_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    drop_last=True,
+)
+
+best_test_acc = 0
+
+for epoch in range(args.num_epochs):
+    train(net, opt, scheduler, train_loader, dev)
+    if (epoch + 1) % 1 == 0:
+        print("Epoch #%d Testing" % epoch)
+        test_acc = evaluate(net, test_loader, dev)
+        if test_acc > best_test_acc:
+            best_test_acc = test_acc
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print("Current test acc: %.5f (best: %.5f)" % (test_acc, best_test_acc))

examples/pytorch/pointcloud/edgeconv/README.md

@@ -0,0 +1,25 @@
+Dynamic EdgeConv
+====
+
+This is a reproduction of the paper [Dynamic Graph CNN for Learning on Point
+Clouds](https://arxiv.org/pdf/1801.07829.pdf).
+
+The reproduced experiment is the 40-class classification on the ModelNet40
+dataset.  The sampled point clouds are identical to that of
+[PointNet](https://github.com/charlesq34/pointnet).
+
+To train and test the model, simply run
+
+```python
+python main.py
+```
+
+The model currently takes 3 minutes to train an epoch on Tesla V100, and an
+additional 17 seconds to run a validation and 20 seconds to run a test.
+
+The best validation performance is 93.5% with a test performance of 91.8%.
+
+## Dependencies
+
+* `h5py`
+* `tqdm`

examples/pytorch/pointcloud/edgeconv/main.py

@@ -0,0 +1,151 @@
+import argparse
+import os
+import urllib
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+
+from dgl.data.utils import download, get_download_dir
+from model import compute_loss, Model
+from modelnet import ModelNet
+from torch.utils.data import DataLoader
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=100)
+parser.add_argument("--num-workers", type=int, default=0)
+parser.add_argument("--batch-size", type=int, default=32)
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+data_filename = "modelnet40-sampled-2048.h5"
+local_path = args.dataset_path or os.path.join(
+    get_download_dir(), data_filename
+)
+
+if not os.path.exists(local_path):
+    download(
+        "https://data.dgl.ai/dataset/modelnet40-sampled-2048.h5", local_path
+    )
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(model, opt, scheduler, train_loader, dev):
+    scheduler.step()
+
+    model.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label in tq:
+            num_examples = label.shape[0]
+            data, label = data.to(dev), label.to(dev).squeeze().long()
+            opt.zero_grad()
+            logits = model(data)
+            loss = compute_loss(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            num_batches += 1
+            count += num_examples
+            loss = loss.item()
+            correct = (preds == label).sum().item()
+            total_loss += loss
+            total_correct += correct
+
+            tq.set_postfix(
+                {
+                    "Loss": "%.5f" % loss,
+                    "AvgLoss": "%.5f" % (total_loss / num_batches),
+                    "Acc": "%.5f" % (correct / num_examples),
+                    "AvgAcc": "%.5f" % (total_correct / count),
+                }
+            )
+
+
+def evaluate(model, test_loader, dev):
+    model.eval()
+
+    total_correct = 0
+    count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label in tq:
+                num_examples = label.shape[0]
+                data, label = data.to(dev), label.to(dev).squeeze().long()
+                logits = model(data)
+                _, preds = logits.max(1)
+
+                correct = (preds == label).sum().item()
+                total_correct += correct
+                count += num_examples
+
+                tq.set_postfix(
+                    {
+                        "Acc": "%.5f" % (correct / num_examples),
+                        "AvgAcc": "%.5f" % (total_correct / count),
+                    }
+                )
+
+    return total_correct / count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model(20, [64, 64, 128, 256], [512, 512, 256], 40)
+model = model.to(dev)
+if args.load_model_path:
+    model.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+opt = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)
+
+scheduler = optim.lr_scheduler.CosineAnnealingLR(
+    opt, args.num_epochs, eta_min=0.001
+)
+
+modelnet = ModelNet(local_path, 1024)
+
+train_loader = CustomDataLoader(modelnet.train())
+valid_loader = CustomDataLoader(modelnet.valid())
+test_loader = CustomDataLoader(modelnet.test())
+
+best_valid_acc = 0
+best_test_acc = 0
+
+for epoch in range(args.num_epochs):
+    print("Epoch #%d Validating" % epoch)
+    valid_acc = evaluate(model, valid_loader, dev)
+    test_acc = evaluate(model, test_loader, dev)
+    if valid_acc > best_valid_acc:
+        best_valid_acc = valid_acc
+        best_test_acc = test_acc
+        if args.save_model_path:
+            torch.save(model.state_dict(), args.save_model_path)
+    print(
+        "Current validation acc: %.5f (best: %.5f), test acc: %.5f (best: %.5f)"
+        % (valid_acc, best_valid_acc, test_acc, best_test_acc)
+    )
+
+    train(model, opt, scheduler, train_loader, dev)

examples/pytorch/pointcloud/edgeconv/model.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dgl.nn.pytorch import EdgeConv, KNNGraph
+
+
+class Model(nn.Module):
+    def __init__(
+        self,
+        k,
+        feature_dims,
+        emb_dims,
+        output_classes,
+        input_dims=3,
+        dropout_prob=0.5,
+    ):
+        super(Model, self).__init__()
+
+        self.nng = KNNGraph(k)
+        self.conv = nn.ModuleList()
+
+        self.num_layers = len(feature_dims)
+        for i in range(self.num_layers):
+            self.conv.append(
+                EdgeConv(
+                    feature_dims[i - 1] if i > 0 else input_dims,
+                    feature_dims[i],
+                    batch_norm=True,
+                )
+            )
+
+        self.proj = nn.Linear(sum(feature_dims), emb_dims[0])
+
+        self.embs = nn.ModuleList()
+        self.bn_embs = nn.ModuleList()
+        self.dropouts = nn.ModuleList()
+
+        self.num_embs = len(emb_dims) - 1
+        for i in range(1, self.num_embs + 1):
+            self.embs.append(
+                nn.Linear(
+                    # * 2 because of concatenation of max- and mean-pooling
+                    emb_dims[i - 1] if i > 1 else (emb_dims[i - 1] * 2),
+                    emb_dims[i],
+                )
+            )
+            self.bn_embs.append(nn.BatchNorm1d(emb_dims[i]))
+            self.dropouts.append(nn.Dropout(dropout_prob))
+
+        self.proj_output = nn.Linear(emb_dims[-1], output_classes)
+
+    def forward(self, x):
+        hs = []
+        batch_size, n_points, x_dims = x.shape
+        h = x
+
+        for i in range(self.num_layers):
+            g = self.nng(h).to(h.device)
+            h = h.view(batch_size * n_points, -1)
+            h = self.conv[i](g, h)
+            h = F.leaky_relu(h, 0.2)
+            h = h.view(batch_size, n_points, -1)
+            hs.append(h)
+
+        h = torch.cat(hs, 2)
+        h = self.proj(h)
+        h_max, _ = torch.max(h, 1)
+        h_avg = torch.mean(h, 1)
+        h = torch.cat([h_max, h_avg], 1)
+
+        for i in range(self.num_embs):
+            h = self.embs[i](h)
+            h = self.bn_embs[i](h)
+            h = F.leaky_relu(h, 0.2)
+            h = self.dropouts[i](h)
+
+        h = self.proj_output(h)
+        return h
+
+
+def compute_loss(logits, y, eps=0.2):
+    num_classes = logits.shape[1]
+    one_hot = torch.zeros_like(logits).scatter_(1, y.view(-1, 1), 1)
+    one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (num_classes - 1)
+    log_prob = F.log_softmax(logits, 1)
+    loss = -(one_hot * log_prob).sum(1).mean()
+    return loss

examples/pytorch/pointcloud/edgeconv/modelnet.py

@@ -0,0 +1,58 @@
+import numpy as np
+from torch.utils.data import Dataset
+
+
+class ModelNet(object):
+    def __init__(self, path, num_points):
+        import h5py
+
+        self.f = h5py.File(path)
+        self.num_points = num_points
+
+        self.n_train = self.f["train/data"].shape[0]
+        self.n_valid = int(self.n_train / 5)
+        self.n_train -= self.n_valid
+        self.n_test = self.f["test/data"].shape[0]
+
+    def train(self):
+        return ModelNetDataset(self, "train")
+
+    def valid(self):
+        return ModelNetDataset(self, "valid")
+
+    def test(self):
+        return ModelNetDataset(self, "test")
+
+
+class ModelNetDataset(Dataset):
+    def __init__(self, modelnet, mode):
+        super(ModelNetDataset, self).__init__()
+        self.num_points = modelnet.num_points
+        self.mode = mode
+
+        if mode == "train":
+            self.data = modelnet.f["train/data"][: modelnet.n_train]
+            self.label = modelnet.f["train/label"][: modelnet.n_train]
+        elif mode == "valid":
+            self.data = modelnet.f["train/data"][modelnet.n_train :]
+            self.label = modelnet.f["train/label"][modelnet.n_train :]
+        elif mode == "test":
+            self.data = modelnet.f["test/data"].value
+            self.label = modelnet.f["test/label"].value
+
+    def translate(self, x, scale=(2 / 3, 3 / 2), shift=(-0.2, 0.2)):
+        xyz1 = np.random.uniform(low=scale[0], high=scale[1], size=[3])
+        xyz2 = np.random.uniform(low=shift[0], high=shift[1], size=[3])
+        x = np.add(np.multiply(x, xyz1), xyz2).astype("float32")
+        return x
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, i):
+        x = self.data[i][: self.num_points]
+        y = self.label[i]
+        if self.mode == "train":
+            x = self.translate(x)
+            np.random.shuffle(x)
+        return x, y

examples/pytorch/pointcloud/pct/ModelNetDataLoader.py

@@ -0,0 +1,128 @@
+import os
+import warnings
+
+import numpy as np
+from torch.utils.data import Dataset
+
+warnings.filterwarnings("ignore")
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def farthest_point_sample(point, npoint):
+    """
+    Farthest point sampler works as follows:
+    1. Initialize the sample set S with a random point
+    2. Pick point P not in S, which maximizes the distance d(P, S)
+    3. Repeat step 2 until |S| = npoint
+
+    Input:
+        xyz: pointcloud data, [N, D]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [npoint, D]
+    """
+    N, D = point.shape
+    xyz = point[:, :3]
+    centroids = np.zeros((npoint,))
+    distance = np.ones((N,)) * 1e10
+    farthest = np.random.randint(0, N)
+    for i in range(npoint):
+        centroids[i] = farthest
+        centroid = xyz[farthest, :]
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+    point = point[centroids.astype(np.int32)]
+    return point
+
+
+class ModelNetDataLoader(Dataset):
+    def __init__(
+        self,
+        root,
+        npoint=1024,
+        split="train",
+        fps=False,
+        normal_channel=True,
+        cache_size=15000,
+    ):
+        """
+        Input:
+            root: the root path to the local data files
+            npoint: number of points from each cloud
+            split: which split of the data, 'train' or 'test'
+            fps: whether to sample points with farthest point sampler
+            normal_channel: whether to use additional channel
+            cache_size: the cache size of in-memory point clouds
+        """
+        self.root = root
+        self.npoints = npoint
+        self.fps = fps
+        self.catfile = os.path.join(self.root, "modelnet40_shape_names.txt")
+
+        self.cat = [line.rstrip() for line in open(self.catfile)]
+        self.classes = dict(zip(self.cat, range(len(self.cat))))
+        self.normal_channel = normal_channel
+
+        shape_ids = {}
+        shape_ids["train"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_train.txt"))
+        ]
+        shape_ids["test"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_test.txt"))
+        ]
+
+        assert split == "train" or split == "test"
+        shape_names = ["_".join(x.split("_")[0:-1]) for x in shape_ids[split]]
+        # list of (shape_name, shape_txt_file_path) tuple
+        self.datapath = [
+            (
+                shape_names[i],
+                os.path.join(self.root, shape_names[i], shape_ids[split][i])
+                + ".txt",
+            )
+            for i in range(len(shape_ids[split]))
+        ]
+        print("The size of %s data is %d" % (split, len(self.datapath)))
+
+        self.cache_size = cache_size
+        self.cache = {}
+
+    def __len__(self):
+        return len(self.datapath)
+
+    def _get_item(self, index):
+        if index in self.cache:
+            point_set, cls = self.cache[index]
+        else:
+            fn = self.datapath[index]
+            cls = self.classes[self.datapath[index][0]]
+            cls = np.array([cls]).astype(np.int32)
+            point_set = np.loadtxt(fn[1], delimiter=",").astype(np.float32)
+            if self.fps:
+                point_set = farthest_point_sample(point_set, self.npoints)
+            else:
+                point_set = point_set[0 : self.npoints, :]
+
+            point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
+
+            if not self.normal_channel:
+                point_set = point_set[:, 0:3]
+
+            if len(self.cache) < self.cache_size:
+                self.cache[index] = (point_set, cls)
+
+        return point_set, cls
+
+    def __getitem__(self, index):
+        return self._get_item(index)

examples/pytorch/pointcloud/pct/README.md

@@ -0,0 +1,29 @@
+PCT
+====
+
+This is a reproduction of the paper: [PCT: Point cloud transformer](http://arxiv.org/abs/2012.09688).
+
+# Performance
+| Task           | Dataset    | Metric   | Score - Paper  | Score - DGL (Adam) | Time(s) - DGL |
+|-----------------|------------|----------|------------------|-------------|-------------------|
+| Classification        | ModelNet40 | Accuracy | 93.2   | 92.1      | 740.0          |
+| Part Segmentation        | ShapeNet   | mIoU     | 86.4            | 85.6       | 390.0         |
+
++ Time(s) are the average training time per epoch, measured on EC2 g4dn.12xlarge instance w/ Tesla T4 GPU.
++ We run the code with the preprocessing used in [PointNet++](../pointnet). We can only get 84.5 for classification if we use the preprocessing described in the paper:
+    > During training, a random translation in [−0.2, 0.2], a random anisotropic scaling in [0.67, 1.5] and a random input dropout were applied to augment the input data.
+
+
+# How to Run
+
+For point cloud classification, run with
+
+```python
+python train_cls.py
+```
+
+For point cloud part-segmentation, run with
+
+```python
+python train_partseg.py
+```

examples/pytorch/pointcloud/pct/ShapeNet.py

@@ -0,0 +1,161 @@
+import json
+import os
+from zipfile import ZipFile
+
+import dgl
+
+import numpy as np
+import tqdm
+from dgl.data.utils import download, get_download_dir
+from scipy.sparse import csr_matrix
+from torch.utils.data import Dataset
+
+
+class ShapeNet(object):
+    def __init__(self, num_points=2048, normal_channel=True):
+        self.num_points = num_points
+        self.normal_channel = normal_channel
+
+        SHAPENET_DOWNLOAD_URL = "https://shapenet.cs.stanford.edu/media/shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        download_path = get_download_dir()
+        data_filename = (
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        )
+        data_path = os.path.join(
+            download_path,
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal",
+        )
+        if not os.path.exists(data_path):
+            local_path = os.path.join(download_path, data_filename)
+            if not os.path.exists(local_path):
+                download(SHAPENET_DOWNLOAD_URL, local_path, verify_ssl=False)
+            with ZipFile(local_path) as z:
+                z.extractall(path=download_path)
+
+        synset_file = "synsetoffset2category.txt"
+        with open(os.path.join(data_path, synset_file)) as f:
+            synset = [t.split("\n")[0].split("\t") for t in f.readlines()]
+        self.synset_dict = {}
+        for syn in synset:
+            self.synset_dict[syn[1]] = syn[0]
+        self.seg_classes = {
+            "Airplane": [0, 1, 2, 3],
+            "Bag": [4, 5],
+            "Cap": [6, 7],
+            "Car": [8, 9, 10, 11],
+            "Chair": [12, 13, 14, 15],
+            "Earphone": [16, 17, 18],
+            "Guitar": [19, 20, 21],
+            "Knife": [22, 23],
+            "Lamp": [24, 25, 26, 27],
+            "Laptop": [28, 29],
+            "Motorbike": [30, 31, 32, 33, 34, 35],
+            "Mug": [36, 37],
+            "Pistol": [38, 39, 40],
+            "Rocket": [41, 42, 43],
+            "Skateboard": [44, 45, 46],
+            "Table": [47, 48, 49],
+        }
+
+        train_split_json = "shuffled_train_file_list.json"
+        val_split_json = "shuffled_val_file_list.json"
+        test_split_json = "shuffled_test_file_list.json"
+        split_path = os.path.join(data_path, "train_test_split")
+        with open(os.path.join(split_path, train_split_json)) as f:
+            tmp = f.read()
+            self.train_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, val_split_json)) as f:
+            tmp = f.read()
+            self.val_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, test_split_json)) as f:
+            tmp = f.read()
+            self.test_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+
+    def train(self):
+        return ShapeNetDataset(
+            self, "train", self.num_points, self.normal_channel
+        )
+
+    def valid(self):
+        return ShapeNetDataset(
+            self, "valid", self.num_points, self.normal_channel
+        )
+
+    def trainval(self):
+        return ShapeNetDataset(
+            self, "trainval", self.num_points, self.normal_channel
+        )
+
+    def test(self):
+        return ShapeNetDataset(
+            self, "test", self.num_points, self.normal_channel
+        )
+
+
+class ShapeNetDataset(Dataset):
+    def __init__(self, shapenet, mode, num_points, normal_channel=True):
+        super(ShapeNetDataset, self).__init__()
+        self.mode = mode
+        self.num_points = num_points
+        if not normal_channel:
+            self.dim = 3
+        else:
+            self.dim = 6
+
+        if mode == "train":
+            self.file_list = shapenet.train_file_list
+        elif mode == "valid":
+            self.file_list = shapenet.val_file_list
+        elif mode == "test":
+            self.file_list = shapenet.test_file_list
+        elif mode == "trainval":
+            self.file_list = shapenet.train_file_list + shapenet.val_file_list
+        else:
+            raise "Not supported `mode`"
+
+        data_list = []
+        label_list = []
+        category_list = []
+        print("Loading data from split " + self.mode)
+        for fn in tqdm.tqdm(self.file_list, ascii=True):
+            with open(fn) as f:
+                data = np.array(
+                    [t.split("\n")[0].split(" ") for t in f.readlines()]
+                ).astype(np.float)
+            data_list.append(data[:, 0 : self.dim])
+            label_list.append(data[:, 6].astype(int))
+            category_list.append(shapenet.synset_dict[fn.split("/")[-2]])
+        self.data = data_list
+        self.label = label_list
+        self.category = category_list
+
+    def translate(self, x, scale=(2 / 3, 3 / 2), shift=(-0.2, 0.2), size=3):
+        xyz1 = np.random.uniform(low=scale[0], high=scale[1], size=[size])
+        xyz2 = np.random.uniform(low=shift[0], high=shift[1], size=[size])
+        x = np.add(np.multiply(x, xyz1), xyz2).astype("float32")
+        return x
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        inds = np.random.choice(
+            self.data[i].shape[0], self.num_points, replace=True
+        )
+        x = self.data[i][inds, : self.dim]
+        y = self.label[i][inds]
+        cat = self.category[i]
+        if self.mode == "train":
+            x = self.translate(x, size=self.dim)
+        x = x.astype(np.float)
+        y = y.astype(int)
+        return x, y, cat

examples/pytorch/pointcloud/pct/helper.py

@@ -0,0 +1,177 @@
+import dgl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.geometry import farthest_point_sampler
+
+"""
+Part of the code are adapted from
+https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+"""
+
+
+def square_distance(src, dst):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src**2, -1).view(B, N, 1)
+    dist += torch.sum(dst**2, -1).view(B, 1, M)
+    return dist
+
+
+def index_points(points, idx):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = (
+        torch.arange(B, dtype=torch.long)
+        .to(device)
+        .view(view_shape)
+        .repeat(repeat_shape)
+    )
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+
+class KNearNeighbors(nn.Module):
+    """
+    Find the k nearest neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(KNearNeighbors, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, pos, centroids):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+        """
+        center_pos = index_points(pos, centroids)
+        sqrdists = square_distance(center_pos, pos)
+        group_idx = sqrdists.argsort(dim=-1)[:, :, : self.n_neighbor]
+        return group_idx
+
+
+class KNNGraphBuilder(nn.Module):
+    """
+    Build NN graph
+    """
+
+    def __init__(self, n_neighbor):
+        super(KNNGraphBuilder, self).__init__()
+        self.n_neighbor = n_neighbor
+        self.knn = KNearNeighbors(n_neighbor)
+
+    def forward(self, pos, centroids, feat=None):
+        dev = pos.device
+        group_idx = self.knn(pos, centroids)
+        B, N, _ = pos.shape
+        glist = []
+        for i in range(B):
+            center = torch.zeros((N)).to(dev)
+            center[centroids[i]] = 1
+            src = group_idx[i].contiguous().view(-1)
+            dst = (
+                centroids[i]
+                .view(-1, 1)
+                .repeat(
+                    1, min(self.n_neighbor, src.shape[0] // centroids.shape[1])
+                )
+                .view(-1)
+            )
+
+            unified = torch.cat([src, dst])
+            uniq, inv_idx = torch.unique(unified, return_inverse=True)
+            src_idx = inv_idx[: src.shape[0]]
+            dst_idx = inv_idx[src.shape[0] :]
+
+            g = dgl.graph((src_idx, dst_idx))
+            g.ndata["pos"] = pos[i][uniq]
+            g.ndata["center"] = center[uniq]
+            if feat is not None:
+                g.ndata["feat"] = feat[i][uniq]
+            glist.append(g)
+        bg = dgl.batch(glist)
+        return bg
+
+
+class KNNMessage(nn.Module):
+    """
+    Compute the input feature from neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(KNNMessage, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, edges):
+        norm = edges.src["feat"] - edges.dst["feat"]
+        if "feat" in edges.src:
+            res = torch.cat([norm, edges.src["feat"]], 1)
+        else:
+            res = norm
+        return {"agg_feat": res}
+
+
+class KNNConv(nn.Module):
+    """
+    Feature aggregation
+    """
+
+    def __init__(self, sizes):
+        super(KNNConv, self).__init__()
+        self.conv = nn.ModuleList()
+        self.bn = nn.ModuleList()
+        for i in range(1, len(sizes)):
+            self.conv.append(nn.Conv2d(sizes[i - 1], sizes[i], 1))
+            self.bn.append(nn.BatchNorm2d(sizes[i]))
+
+    def forward(self, nodes):
+        shape = nodes.mailbox["agg_feat"].shape
+        h = (
+            nodes.mailbox["agg_feat"]
+            .view(shape[0], -1, shape[1], shape[2])
+            .permute(0, 3, 2, 1)
+        )
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h, 2)[0]
+        feat_dim = h.shape[1]
+        h = h.permute(0, 2, 1).reshape(-1, feat_dim)
+        return {"new_feat": h}
+
+
+class TransitionDown(nn.Module):
+    """
+    The Transition Down Module
+    """
+
+    def __init__(self, in_channels, out_channels, n_neighbor=64):
+        super(TransitionDown, self).__init__()
+        self.frnn_graph = KNNGraphBuilder(n_neighbor)
+        self.message = KNNMessage(n_neighbor)
+        self.conv = KNNConv([in_channels, out_channels, out_channels])
+
+    def forward(self, pos, feat, n_point):
+        batch_size = pos.shape[0]
+        centroids = farthest_point_sampler(pos, n_point)
+        g = self.frnn_graph(pos, centroids, feat)
+        g.update_all(self.message, self.conv)
+
+        mask = g.ndata["center"] == 1
+        pos_dim = g.ndata["pos"].shape[-1]
+        feat_dim = g.ndata["new_feat"].shape[-1]
+        pos_res = g.ndata["pos"][mask].view(batch_size, -1, pos_dim)
+        feat_res = g.ndata["new_feat"][mask].view(batch_size, -1, feat_dim)
+        return pos_res, feat_res

examples/pytorch/pointcloud/pct/pct.py

@@ -0,0 +1,226 @@
+import torch
+from helper import TransitionDown
+from torch import nn
+
+"""
+Part of the code are adapted from
+https://github.com/MenghaoGuo/PCT
+"""
+
+
+class PCTPositionEmbedding(nn.Module):
+    def __init__(self, channels=256):
+        super(PCTPositionEmbedding, self).__init__()
+        self.conv1 = nn.Conv1d(channels, channels, kernel_size=1, bias=False)
+        self.conv_pos = nn.Conv1d(3, channels, kernel_size=1, bias=False)
+
+        self.bn1 = nn.BatchNorm1d(channels)
+
+        self.sa1 = SALayerCLS(channels)
+        self.sa2 = SALayerCLS(channels)
+        self.sa3 = SALayerCLS(channels)
+        self.sa4 = SALayerCLS(channels)
+
+        self.relu = nn.ReLU()
+
+    def forward(self, x, xyz):
+        # add position embedding
+        xyz = xyz.permute(0, 2, 1)
+        xyz = self.conv_pos(xyz)
+
+        x = self.relu(self.bn1(self.conv1(x)))  # B, D, N
+
+        x1 = self.sa1(x, xyz)
+        x2 = self.sa2(x1, xyz)
+        x3 = self.sa3(x2, xyz)
+        x4 = self.sa4(x3, xyz)
+
+        x = torch.cat((x1, x2, x3, x4), dim=1)
+
+        return x
+
+
+class SALayerCLS(nn.Module):
+    def __init__(self, channels):
+        super(SALayerCLS, self).__init__()
+        self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
+        self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
+        self.q_conv.weight = self.k_conv.weight
+        self.v_conv = nn.Conv1d(channels, channels, 1)
+        self.trans_conv = nn.Conv1d(channels, channels, 1)
+        self.after_norm = nn.BatchNorm1d(channels)
+        self.act = nn.ReLU()
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, xyz):
+        x = x + xyz
+        x_q = self.q_conv(x).permute(0, 2, 1)  # b, n, c
+        x_k = self.k_conv(x)  # b, c, n
+        x_v = self.v_conv(x)
+        energy = torch.bmm(x_q, x_k)  # b, n, n
+        attention = self.softmax(energy)
+        attention = attention / (1e-9 + attention.sum(dim=1, keepdims=True))
+        x_r = torch.bmm(x_v, attention)  # b, c, n
+        x_r = self.act(self.after_norm(self.trans_conv(x - x_r)))
+        x = x + x_r
+        return x
+
+
+class SALayerSeg(nn.Module):
+    def __init__(self, channels):
+        super(SALayerSeg, self).__init__()
+        self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
+        self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
+        self.q_conv.weight = self.k_conv.weight
+        self.v_conv = nn.Conv1d(channels, channels, 1)
+        self.trans_conv = nn.Conv1d(channels, channels, 1)
+        self.after_norm = nn.BatchNorm1d(channels)
+        self.act = nn.ReLU()
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        x_q = self.q_conv(x).permute(0, 2, 1)  # b, n, c
+        x_k = self.k_conv(x)  # b, c, n
+        x_v = self.v_conv(x)
+        energy = torch.bmm(x_q, x_k)  # b, n, n
+        attention = self.softmax(energy)
+        attention = attention / (1e-9 + attention.sum(dim=1, keepdims=True))
+        x_r = torch.bmm(x_v, attention)  # b, c, n
+        x_r = self.act(self.after_norm(self.trans_conv(x - x_r)))
+        x = x + x_r
+        return x
+
+
+class PointTransformerCLS(nn.Module):
+    def __init__(self, output_channels=40):
+        super(PointTransformerCLS, self).__init__()
+        self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False)
+        self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(64)
+        self.g_op0 = TransitionDown(
+            in_channels=128, out_channels=128, n_neighbor=32
+        )
+        self.g_op1 = TransitionDown(
+            in_channels=256, out_channels=256, n_neighbor=32
+        )
+
+        self.pt_last = PCTPositionEmbedding()
+
+        self.relu = nn.ReLU()
+        self.conv_fuse = nn.Sequential(
+            nn.Conv1d(1280, 1024, kernel_size=1, bias=False),
+            nn.BatchNorm1d(1024),
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+
+        self.linear1 = nn.Linear(1024, 512, bias=False)
+        self.bn6 = nn.BatchNorm1d(512)
+        self.dp1 = nn.Dropout(p=0.5)
+        self.linear2 = nn.Linear(512, 256)
+        self.bn7 = nn.BatchNorm1d(256)
+        self.dp2 = nn.Dropout(p=0.5)
+        self.linear3 = nn.Linear(256, output_channels)
+
+    def forward(self, x):
+        xyz = x[..., :3]
+        x = x[..., 3:].permute(0, 2, 1)
+        batch_size, _, _ = x.size()
+        x = self.relu(self.bn1(self.conv1(x)))  # B, D, N
+        x = self.relu(self.bn2(self.conv2(x)))  # B, D, N
+        x = x.permute(0, 2, 1)
+
+        new_xyz, feature_0 = self.g_op0(xyz, x, n_point=512)
+        new_xyz, feature_1 = self.g_op1(new_xyz, feature_0, n_point=256)
+
+        # add position embedding on each layer
+        x = self.pt_last(feature_1, new_xyz)
+
+        x = torch.cat([x, feature_1], dim=1)
+        x = self.conv_fuse(x)
+        x, _ = torch.max(x, 2)
+        x = x.view(batch_size, -1)
+
+        x = self.relu(self.bn6(self.linear1(x)))
+        x = self.dp1(x)
+        x = self.relu(self.bn7(self.linear2(x)))
+        x = self.dp2(x)
+        x = self.linear3(x)
+
+        return x
+
+
+class PointTransformerSeg(nn.Module):
+    def __init__(self, part_num=50):
+        super(PointTransformerSeg, self).__init__()
+        self.part_num = part_num
+        self.conv1 = nn.Conv1d(3, 128, kernel_size=1, bias=False)
+        self.conv2 = nn.Conv1d(128, 128, kernel_size=1, bias=False)
+
+        self.bn1 = nn.BatchNorm1d(128)
+        self.bn2 = nn.BatchNorm1d(128)
+
+        self.sa1 = SALayerSeg(128)
+        self.sa2 = SALayerSeg(128)
+        self.sa3 = SALayerSeg(128)
+        self.sa4 = SALayerSeg(128)
+
+        self.conv_fuse = nn.Sequential(
+            nn.Conv1d(512, 1024, kernel_size=1, bias=False),
+            nn.BatchNorm1d(1024),
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+
+        self.label_conv = nn.Sequential(
+            nn.Conv1d(16, 64, kernel_size=1, bias=False),
+            nn.BatchNorm1d(64),
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+
+        self.convs1 = nn.Conv1d(1024 * 3 + 64, 512, 1)
+        self.dp1 = nn.Dropout(0.5)
+        self.convs2 = nn.Conv1d(512, 256, 1)
+        self.convs3 = nn.Conv1d(256, self.part_num, 1)
+        self.bns1 = nn.BatchNorm1d(512)
+        self.bns2 = nn.BatchNorm1d(256)
+
+        self.relu = nn.ReLU()
+
+    def forward(self, x, cls_label):
+        x = x.permute(0, 2, 1)
+        batch_size, _, N = x.size()
+        x = self.relu(self.bn1(self.conv1(x)))  # B, D, N
+        x = self.relu(self.bn2(self.conv2(x)))
+        x1 = self.sa1(x)
+        x2 = self.sa2(x1)
+        x3 = self.sa3(x2)
+        x4 = self.sa4(x3)
+        x = torch.cat((x1, x2, x3, x4), dim=1)
+        x = self.conv_fuse(x)
+        x_max, _ = torch.max(x, 2)
+        x_avg = torch.mean(x, 2)
+        x_max_feature = x_max.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N)
+        x_avg_feature = x_avg.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N)
+        cls_label_feature = self.label_conv(cls_label).repeat(1, 1, N)
+        x_global_feature = torch.cat(
+            (x_max_feature, x_avg_feature, cls_label_feature), 1
+        )
+        x = torch.cat((x, x_global_feature), 1)
+        x = self.relu(self.bns1(self.convs1(x)))
+        x = self.dp1(x)
+        x = self.relu(self.bns2(self.convs2(x)))
+        x = self.convs3(x)
+        return x
+
+
+class PartSegLoss(nn.Module):
+    def __init__(self, eps=0.2):
+        super(PartSegLoss, self).__init__()
+        self.eps = eps
+        self.loss = nn.CrossEntropyLoss()
+
+    def forward(self, logits, y):
+        num_classes = logits.shape[1]
+        logits = logits.permute(0, 2, 1).contiguous().view(-1, num_classes)
+        loss = self.loss(logits, y)
+        return loss

examples/pytorch/pointcloud/pct/provider.py

@@ -0,0 +1,312 @@
+"""
+Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/master/provider.py
+"""
+import numpy as np
+
+
+def normalize_data(batch_data):
+    """Normalize the batch data, use coordinates of the block centered at origin,
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    B, N, C = batch_data.shape
+    normal_data = np.zeros((B, N, C))
+    for b in range(B):
+        pc = batch_data[b]
+        centroid = np.mean(pc, axis=0)
+        pc = pc - centroid
+        m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+        pc = pc / m
+        normal_data[b] = pc
+    return normal_data
+
+
+def shuffle_data(data, labels):
+    """Shuffle data and labels.
+    Input:
+      data: B,N,... numpy array
+      label: B,... numpy array
+    Return:
+      shuffled data, label and shuffle indices
+    """
+    idx = np.arange(len(labels))
+    np.random.shuffle(idx)
+    return data[idx, ...], labels[idx], idx
+
+
+def shuffle_points(batch_data):
+    """Shuffle orders of points in each point cloud -- changes FPS behavior.
+    Use the same shuffling idx for the entire batch.
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    idx = np.arange(batch_data.shape[1])
+    np.random.shuffle(idx)
+    return batch_data[:, idx, :]
+
+
+def rotate_point_cloud(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_z(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, sinval, 0], [-sinval, cosval, 0], [0, 0, 1]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_with_normal(batch_xyz_normal):
+    """Randomly rotate XYZ, normal point cloud.
+    Input:
+        batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
+    Output:
+        B,N,6, rotated XYZ, normal point cloud
+    """
+    for k in range(batch_xyz_normal.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_xyz_normal[k, :, 0:3]
+        shape_normal = batch_xyz_normal[k, :, 3:6]
+        batch_xyz_normal[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        batch_xyz_normal[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return batch_xyz_normal
+
+
+def rotate_perturbation_point_cloud_with_normal(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx6 array, original batch of point clouds and point normals
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
+        rotated_data[k, :, 3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx6 array, original batch of point clouds with normal
+      scalar, angle of rotation
+    Return:
+      BxNx6 array, rotated batch of point clouds iwth normal
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        rotated_data[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_perturbation_point_cloud(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
+    """Randomly jitter points. jittering is per point.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, jittered batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    assert clip > 0
+    jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1 * clip, clip)
+    jittered_data += batch_data
+    return jittered_data
+
+
+def shift_point_cloud(batch_data, shift_range=0.1):
+    """Randomly shift point cloud. Shift is per point cloud.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, shifted batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    shifts = np.random.uniform(-shift_range, shift_range, (B, 3))
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] += shifts[batch_index, :]
+    return batch_data
+
+
+def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25):
+    """Randomly scale the point cloud. Scale is per point cloud.
+    Input:
+        BxNx3 array, original batch of point clouds
+    Return:
+        BxNx3 array, scaled batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    scales = np.random.uniform(scale_low, scale_high, B)
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] *= scales[batch_index]
+    return batch_data
+
+
+def random_point_dropout(batch_pc, max_dropout_ratio=0.875):
+    """batch_pc: BxNx3"""
+    for b in range(batch_pc.shape[0]):
+        dropout_ratio = np.random.random() * max_dropout_ratio  # 0~0.875
+        drop_idx = np.where(
+            np.random.random((batch_pc.shape[1])) <= dropout_ratio
+        )[0]
+        if len(drop_idx) > 0:
+            dropout_ratio = (
+                np.random.random() * max_dropout_ratio
+            )  # 0~0.875 # not need
+            batch_pc[b, drop_idx, :] = batch_pc[
+                b, 0, :
+            ]  # set to the first point
+    return batch_pc

examples/pytorch/pointcloud/pct/train_cls.py

@@ -0,0 +1,183 @@
+import argparse
+import os
+import time
+from functools import partial
+
+import provider
+import torch
+import torch.nn as nn
+import tqdm
+
+from dgl.data.utils import download, get_download_dir
+from ModelNetDataLoader import ModelNetDataLoader
+from pct import PointTransformerCLS
+from torch.utils.data import DataLoader
+
+torch.backends.cudnn.enabled = False
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=250)
+parser.add_argument("--num-workers", type=int, default=8)
+parser.add_argument("--batch-size", type=int, default=32)
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+data_filename = "modelnet40_normal_resampled.zip"
+download_path = os.path.join(get_download_dir(), data_filename)
+local_path = args.dataset_path or os.path.join(
+    get_download_dir(), "modelnet40_normal_resampled"
+)
+
+if not os.path.exists(local_path):
+    download(
+        "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip",
+        download_path,
+        verify_ssl=False,
+    )
+    from zipfile import ZipFile
+
+    with ZipFile(download_path) as z:
+        z.extractall(path=get_download_dir())
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    loss_f = nn.CrossEntropyLoss()
+    start_time = time.time()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label in tq:
+            data = data.data.numpy()
+            data = provider.random_point_dropout(data)
+            data[:, :, 0:3] = provider.random_scale_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.jitter_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.shift_point_cloud(data[:, :, 0:3])
+            data = torch.tensor(data)
+            label = label[:, 0]
+
+            num_examples = label.shape[0]
+            data, label = data.to(dev), label.to(dev).squeeze().long()
+            opt.zero_grad()
+            logits = net(data)
+            loss = loss_f(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            num_batches += 1
+            count += num_examples
+            loss = loss.item()
+            correct = (preds == label).sum().item()
+            total_loss += loss
+            total_correct += correct
+
+            tq.set_postfix(
+                {
+                    "AvgLoss": "%.5f" % (total_loss / num_batches),
+                    "AvgAcc": "%.5f" % (total_correct / count),
+                }
+            )
+    print(
+        "[Train] AvgLoss: {:.5}, AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_loss / num_batches,
+            total_correct / count,
+            time.time() - start_time,
+        )
+    )
+    scheduler.step()
+
+
+def evaluate(net, test_loader, dev):
+    net.eval()
+
+    total_correct = 0
+    count = 0
+    start_time = time.time()
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label in tq:
+                label = label[:, 0]
+                num_examples = label.shape[0]
+                data, label = data.to(dev), label.to(dev).squeeze().long()
+                logits = net(data)
+                _, preds = logits.max(1)
+
+                correct = (preds == label).sum().item()
+                total_correct += correct
+                count += num_examples
+
+                tq.set_postfix({"AvgAcc": "%.5f" % (total_correct / count)})
+    print(
+        "[Test]  AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_correct / count, time.time() - start_time
+        )
+    )
+    return total_correct / count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net = PointTransformerCLS()
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+
+opt = torch.optim.SGD(
+    net.parameters(), lr=0.01, weight_decay=1e-4, momentum=0.9
+)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+    opt, T_max=args.num_epochs
+)
+
+train_dataset = ModelNetDataLoader(local_path, 1024, split="train")
+test_dataset = ModelNetDataLoader(local_path, 1024, split="test")
+train_loader = torch.utils.data.DataLoader(
+    train_dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    num_workers=num_workers,
+    drop_last=True,
+)
+test_loader = torch.utils.data.DataLoader(
+    test_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    drop_last=True,
+)
+
+best_test_acc = 0
+
+for epoch in range(args.num_epochs):
+    print("Epoch #{}: ".format(epoch))
+    train(net, opt, scheduler, train_loader, dev)
+    if (epoch + 1) % 1 == 0:
+        test_acc = evaluate(net, test_loader, dev)
+        if test_acc > best_test_acc:
+            best_test_acc = test_acc
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print("Current test acc: %.5f (best: %.5f)" % (test_acc, best_test_acc))
+    print()

examples/pytorch/pointcloud/pct/train_partseg.py

@@ -0,0 +1,314 @@
+import argparse
+import time
+from functools import partial
+
+import dgl
+
+import numpy as np
+import provider
+import torch
+import torch.optim as optim
+import tqdm
+from pct import PartSegLoss, PointTransformerSeg
+from ShapeNet import ShapeNet
+from torch.utils.data import DataLoader
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=500)
+parser.add_argument("--num-workers", type=int, default=8)
+parser.add_argument("--batch-size", type=int, default=16)
+parser.add_argument("--tensorboard", action="store_true")
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+
+def collate(samples):
+    graphs, cat = map(list, zip(*samples))
+    return dgl.batch(graphs), cat
+
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    start = time.time()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label, cat in tq:
+            num_examples = data.shape[0]
+            data = data.to(dev, dtype=torch.float)
+            label = label.to(dev, dtype=torch.long).view(-1)
+            opt.zero_grad()
+            cat_ind = [category_list.index(c) for c in cat]
+            # An one-hot encoding for the object category
+            cat_tensor = torch.tensor(eye_mat[cat_ind]).to(
+                dev, dtype=torch.float
+            )
+            cat_tensor = cat_tensor.view(num_examples, 16, 1)
+            logits = net(data, cat_tensor)
+            loss = L(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            count += num_examples * 2048
+            loss = loss.item()
+            total_loss += loss
+            num_batches += 1
+            correct = (preds.view(-1) == label).sum().item()
+            total_correct += correct
+
+            AvgLoss = total_loss / num_batches
+            AvgAcc = total_correct / count
+
+            tq.set_postfix(
+                {"AvgLoss": "%.5f" % AvgLoss, "AvgAcc": "%.5f" % AvgAcc}
+            )
+    scheduler.step()
+    end = time.time()
+    print(
+        "[Train] AvgLoss: {:.5}, AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_loss / num_batches, total_correct / count, end - start
+        )
+    )
+    return data, preds, AvgLoss, AvgAcc, end - start
+
+
+def mIoU(preds, label, cat, cat_miou, seg_classes):
+    for i in range(preds.shape[0]):
+        shape_iou = 0
+        n = len(seg_classes[cat[i]])
+        for cls in seg_classes[cat[i]]:
+            pred_set = set(np.where(preds[i, :] == cls)[0])
+            label_set = set(np.where(label[i, :] == cls)[0])
+            union = len(pred_set.union(label_set))
+            inter = len(pred_set.intersection(label_set))
+            if union == 0:
+                shape_iou += 1
+            else:
+                shape_iou += inter / union
+        shape_iou /= n
+        cat_miou[cat[i]][0] += shape_iou
+        cat_miou[cat[i]][1] += 1
+
+    return cat_miou
+
+
+def evaluate(net, test_loader, dev, per_cat_verbose=False):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.eval()
+
+    cat_miou = {}
+    for k in shapenet.seg_classes.keys():
+        cat_miou[k] = [0, 0]
+    miou = 0
+    count = 0
+    per_cat_miou = 0
+    per_cat_count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label, cat in tq:
+                num_examples = data.shape[0]
+                data = data.to(dev, dtype=torch.float)
+                label = label.to(dev, dtype=torch.long)
+                cat_ind = [category_list.index(c) for c in cat]
+                cat_tensor = torch.tensor(eye_mat[cat_ind]).to(
+                    dev, dtype=torch.float
+                )
+                cat_tensor = cat_tensor.view(num_examples, 16, 1)
+                logits = net(data, cat_tensor)
+                _, preds = logits.max(1)
+
+                cat_miou = mIoU(
+                    preds.cpu().numpy(),
+                    label.view(num_examples, -1).cpu().numpy(),
+                    cat,
+                    cat_miou,
+                    shapenet.seg_classes,
+                )
+                for _, v in cat_miou.items():
+                    if v[1] > 0:
+                        miou += v[0]
+                        count += v[1]
+                        per_cat_miou += v[0] / v[1]
+                        per_cat_count += 1
+                tq.set_postfix(
+                    {
+                        "mIoU": "%.5f" % (miou / count),
+                        "per Category mIoU": "%.5f"
+                        % (per_cat_miou / per_cat_count),
+                    }
+                )
+    print(
+        "[Test] mIoU: %.5f, per Category mIoU: %.5f"
+        % (miou / count, per_cat_miou / per_cat_count)
+    )
+    if per_cat_verbose:
+        print("-" * 60)
+        print("Per-Category mIoU:")
+        for k, v in cat_miou.items():
+            if v[1] > 0:
+                print("%s mIoU=%.5f" % (k, v[0] / v[1]))
+            else:
+                print("%s mIoU=%.5f" % (k, 1))
+        print("-" * 60)
+    return miou / count, per_cat_miou / per_cat_count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net = PointTransformerSeg()
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+opt = torch.optim.SGD(
+    net.parameters(), lr=0.01, weight_decay=1e-4, momentum=0.9
+)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+    opt, T_max=args.num_epochs
+)
+
+L = PartSegLoss()
+
+shapenet = ShapeNet(2048, normal_channel=False)
+
+train_loader = CustomDataLoader(shapenet.trainval())
+test_loader = CustomDataLoader(shapenet.test())
+
+# Tensorboard
+if args.tensorboard:
+    import torchvision
+    from torch.utils.tensorboard import SummaryWriter
+    from torchvision import datasets, transforms
+
+    writer = SummaryWriter()
+# Select 50 distinct colors for different parts
+color_map = torch.tensor(
+    [
+        [47, 79, 79],
+        [139, 69, 19],
+        [112, 128, 144],
+        [85, 107, 47],
+        [139, 0, 0],
+        [128, 128, 0],
+        [72, 61, 139],
+        [0, 128, 0],
+        [188, 143, 143],
+        [60, 179, 113],
+        [205, 133, 63],
+        [0, 139, 139],
+        [70, 130, 180],
+        [205, 92, 92],
+        [154, 205, 50],
+        [0, 0, 139],
+        [50, 205, 50],
+        [250, 250, 250],
+        [218, 165, 32],
+        [139, 0, 139],
+        [10, 10, 10],
+        [176, 48, 96],
+        [72, 209, 204],
+        [153, 50, 204],
+        [255, 69, 0],
+        [255, 145, 0],
+        [0, 0, 205],
+        [255, 255, 0],
+        [0, 255, 0],
+        [233, 150, 122],
+        [220, 20, 60],
+        [0, 191, 255],
+        [160, 32, 240],
+        [192, 192, 192],
+        [173, 255, 47],
+        [218, 112, 214],
+        [216, 191, 216],
+        [255, 127, 80],
+        [255, 0, 255],
+        [100, 149, 237],
+        [128, 128, 128],
+        [221, 160, 221],
+        [144, 238, 144],
+        [123, 104, 238],
+        [255, 160, 122],
+        [175, 238, 238],
+        [238, 130, 238],
+        [127, 255, 212],
+        [255, 218, 185],
+        [255, 105, 180],
+    ]
+)
+# paint each point according to its pred
+
+
+def paint(batched_points):
+    B, N = batched_points.shape
+    colored = color_map[batched_points].squeeze(2)
+    return colored
+
+
+best_test_miou = 0
+best_test_per_cat_miou = 0
+
+for epoch in range(args.num_epochs):
+    print("Epoch #{}: ".format(epoch))
+    data, preds, AvgLoss, AvgAcc, training_time = train(
+        net, opt, scheduler, train_loader, dev
+    )
+    if (epoch + 1) % 5 == 0 or epoch == 0:
+        test_miou, test_per_cat_miou = evaluate(net, test_loader, dev, True)
+        if test_miou > best_test_miou:
+            best_test_miou = test_miou
+            best_test_per_cat_miou = test_per_cat_miou
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print(
+            "Current test mIoU: %.5f (best: %.5f), per-Category mIoU: %.5f (best: %.5f)"
+            % (
+                test_miou,
+                best_test_miou,
+                test_per_cat_miou,
+                best_test_per_cat_miou,
+            )
+        )
+    # Tensorboard
+    if args.tensorboard:
+        colored = paint(preds)
+        writer.add_mesh(
+            "data", vertices=data, colors=colored, global_step=epoch
+        )
+        writer.add_scalar(
+            "training time for one epoch", training_time, global_step=epoch
+        )
+        writer.add_scalar("AvgLoss", AvgLoss, global_step=epoch)
+        writer.add_scalar("AvgAcc", AvgAcc, global_step=epoch)
+        if (epoch + 1) % 5 == 0:
+            writer.add_scalar("test mIoU", test_miou, global_step=epoch)
+            writer.add_scalar(
+                "best test mIoU", best_test_miou, global_step=epoch
+            )
+    print()

examples/pytorch/pointcloud/point_transformer/ModelNetDataLoader.py

@@ -0,0 +1,128 @@
+import os
+import warnings
+
+import numpy as np
+from torch.utils.data import Dataset
+
+warnings.filterwarnings("ignore")
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def farthest_point_sample(point, npoint):
+    """
+    Farthest point sampler works as follows:
+    1. Initialize the sample set S with a random point
+    2. Pick point P not in S, which maximizes the distance d(P, S)
+    3. Repeat step 2 until |S| = npoint
+
+    Input:
+        xyz: pointcloud data, [N, D]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [npoint, D]
+    """
+    N, D = point.shape
+    xyz = point[:, :3]
+    centroids = np.zeros((npoint,))
+    distance = np.ones((N,)) * 1e10
+    farthest = np.random.randint(0, N)
+    for i in range(npoint):
+        centroids[i] = farthest
+        centroid = xyz[farthest, :]
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+    point = point[centroids.astype(np.int32)]
+    return point
+
+
+class ModelNetDataLoader(Dataset):
+    def __init__(
+        self,
+        root,
+        npoint=1024,
+        split="train",
+        fps=False,
+        normal_channel=True,
+        cache_size=15000,
+    ):
+        """
+        Input:
+            root: the root path to the local data files
+            npoint: number of points from each cloud
+            split: which split of the data, 'train' or 'test'
+            fps: whether to sample points with farthest point sampler
+            normal_channel: whether to use additional channel
+            cache_size: the cache size of in-memory point clouds
+        """
+        self.root = root
+        self.npoints = npoint
+        self.fps = fps
+        self.catfile = os.path.join(self.root, "modelnet40_shape_names.txt")
+
+        self.cat = [line.rstrip() for line in open(self.catfile)]
+        self.classes = dict(zip(self.cat, range(len(self.cat))))
+        self.normal_channel = normal_channel
+
+        shape_ids = {}
+        shape_ids["train"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_train.txt"))
+        ]
+        shape_ids["test"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_test.txt"))
+        ]
+
+        assert split == "train" or split == "test"
+        shape_names = ["_".join(x.split("_")[0:-1]) for x in shape_ids[split]]
+        # list of (shape_name, shape_txt_file_path) tuple
+        self.datapath = [
+            (
+                shape_names[i],
+                os.path.join(self.root, shape_names[i], shape_ids[split][i])
+                + ".txt",
+            )
+            for i in range(len(shape_ids[split]))
+        ]
+        print("The size of %s data is %d" % (split, len(self.datapath)))
+
+        self.cache_size = cache_size
+        self.cache = {}
+
+    def __len__(self):
+        return len(self.datapath)
+
+    def _get_item(self, index):
+        if index in self.cache:
+            point_set, cls = self.cache[index]
+        else:
+            fn = self.datapath[index]
+            cls = self.classes[self.datapath[index][0]]
+            cls = np.array([cls]).astype(np.int32)
+            point_set = np.loadtxt(fn[1], delimiter=",").astype(np.float32)
+            if self.fps:
+                point_set = farthest_point_sample(point_set, self.npoints)
+            else:
+                point_set = point_set[0 : self.npoints, :]
+
+            point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
+
+            if not self.normal_channel:
+                point_set = point_set[:, 0:3]
+
+            if len(self.cache) < self.cache_size:
+                self.cache[index] = (point_set, cls)
+
+        return point_set, cls
+
+    def __getitem__(self, index):
+        return self._get_item(index)

examples/pytorch/pointcloud/point_transformer/README.md

@@ -0,0 +1,29 @@
+Point Transformer
+====
+
+> This model is implemented on August 27, 2021 when there is no official code released.    
+Thus we implemented this model based on the code from <https://github.com/qq456cvb/Point-Transformers>.
+
+This is a reproduction of the paper: [Point Transformer](http://arxiv.org/abs/2012.09164).
+
+# Performance
+| Task           | Dataset    | Metric   | Score - Paper  | Score - DGL (Adam) | Score - DGL (SGD) | Time(s) - DGL |
+|-----------------|------------|----------|------------------|-------------|-------------|-------------------|
+| Classification        | ModelNet40 | Accuracy | 93.7   | 92.0        |  91.5        | 117.0          |
+| Part Segmentation        | ShapeNet   | mIoU     | 86.6            | 84.3        |  85.1        | 260.0         |
+
++ Time(s) are the average training time per epoch, measured on EC2 p3.8xlarge instance w/ Tesla V100 GPU.
+
+# How to Run
+
+For point cloud classification, run with
+
+```python
+python train_cls.py --opt [sgd/adam]
+```
+
+For point cloud part-segmentation, run with
+
+```python
+python train_partseg.py --opt [sgd/adam]
+```

examples/pytorch/pointcloud/point_transformer/ShapeNet.py

@@ -0,0 +1,161 @@
+import json
+import os
+from zipfile import ZipFile
+
+import dgl
+
+import numpy as np
+import tqdm
+from dgl.data.utils import download, get_download_dir
+from scipy.sparse import csr_matrix
+from torch.utils.data import Dataset
+
+
+class ShapeNet(object):
+    def __init__(self, num_points=2048, normal_channel=True):
+        self.num_points = num_points
+        self.normal_channel = normal_channel
+
+        SHAPENET_DOWNLOAD_URL = "https://shapenet.cs.stanford.edu/media/shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        download_path = get_download_dir()
+        data_filename = (
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        )
+        data_path = os.path.join(
+            download_path,
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal",
+        )
+        if not os.path.exists(data_path):
+            local_path = os.path.join(download_path, data_filename)
+            if not os.path.exists(local_path):
+                download(SHAPENET_DOWNLOAD_URL, local_path, verify_ssl=False)
+            with ZipFile(local_path) as z:
+                z.extractall(path=download_path)
+
+        synset_file = "synsetoffset2category.txt"
+        with open(os.path.join(data_path, synset_file)) as f:
+            synset = [t.split("\n")[0].split("\t") for t in f.readlines()]
+        self.synset_dict = {}
+        for syn in synset:
+            self.synset_dict[syn[1]] = syn[0]
+        self.seg_classes = {
+            "Airplane": [0, 1, 2, 3],
+            "Bag": [4, 5],
+            "Cap": [6, 7],
+            "Car": [8, 9, 10, 11],
+            "Chair": [12, 13, 14, 15],
+            "Earphone": [16, 17, 18],
+            "Guitar": [19, 20, 21],
+            "Knife": [22, 23],
+            "Lamp": [24, 25, 26, 27],
+            "Laptop": [28, 29],
+            "Motorbike": [30, 31, 32, 33, 34, 35],
+            "Mug": [36, 37],
+            "Pistol": [38, 39, 40],
+            "Rocket": [41, 42, 43],
+            "Skateboard": [44, 45, 46],
+            "Table": [47, 48, 49],
+        }
+
+        train_split_json = "shuffled_train_file_list.json"
+        val_split_json = "shuffled_val_file_list.json"
+        test_split_json = "shuffled_test_file_list.json"
+        split_path = os.path.join(data_path, "train_test_split")
+        with open(os.path.join(split_path, train_split_json)) as f:
+            tmp = f.read()
+            self.train_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, val_split_json)) as f:
+            tmp = f.read()
+            self.val_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, test_split_json)) as f:
+            tmp = f.read()
+            self.test_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+
+    def train(self):
+        return ShapeNetDataset(
+            self, "train", self.num_points, self.normal_channel
+        )
+
+    def valid(self):
+        return ShapeNetDataset(
+            self, "valid", self.num_points, self.normal_channel
+        )
+
+    def trainval(self):
+        return ShapeNetDataset(
+            self, "trainval", self.num_points, self.normal_channel
+        )
+
+    def test(self):
+        return ShapeNetDataset(
+            self, "test", self.num_points, self.normal_channel
+        )
+
+
+class ShapeNetDataset(Dataset):
+    def __init__(self, shapenet, mode, num_points, normal_channel=True):
+        super(ShapeNetDataset, self).__init__()
+        self.mode = mode
+        self.num_points = num_points
+        if not normal_channel:
+            self.dim = 3
+        else:
+            self.dim = 6
+
+        if mode == "train":
+            self.file_list = shapenet.train_file_list
+        elif mode == "valid":
+            self.file_list = shapenet.val_file_list
+        elif mode == "test":
+            self.file_list = shapenet.test_file_list
+        elif mode == "trainval":
+            self.file_list = shapenet.train_file_list + shapenet.val_file_list
+        else:
+            raise "Not supported `mode`"
+
+        data_list = []
+        label_list = []
+        category_list = []
+        print("Loading data from split " + self.mode)
+        for fn in tqdm.tqdm(self.file_list, ascii=True):
+            with open(fn) as f:
+                data = np.array(
+                    [t.split("\n")[0].split(" ") for t in f.readlines()]
+                ).astype(float)
+            data_list.append(data[:, 0 : self.dim])
+            label_list.append(data[:, 6].astype(int))
+            category_list.append(shapenet.synset_dict[fn.split("/")[-2]])
+        self.data = data_list
+        self.label = label_list
+        self.category = category_list
+
+    def translate(self, x, scale=(2 / 3, 3 / 2), shift=(-0.2, 0.2), size=3):
+        xyz1 = np.random.uniform(low=scale[0], high=scale[1], size=[size])
+        xyz2 = np.random.uniform(low=shift[0], high=shift[1], size=[size])
+        x = np.add(np.multiply(x, xyz1), xyz2).astype("float32")
+        return x
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        inds = np.random.choice(
+            self.data[i].shape[0], self.num_points, replace=True
+        )
+        x = self.data[i][inds, : self.dim]
+        y = self.label[i][inds]
+        cat = self.category[i]
+        if self.mode == "train":
+            x = self.translate(x, size=self.dim)
+        x = x.astype(float)
+        y = y.astype(int)
+        return x, y, cat

examples/pytorch/pointcloud/point_transformer/helper.py

@@ -0,0 +1,293 @@
+import dgl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.geometry import farthest_point_sampler
+
+"""
+Part of the code are adapted from
+https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+"""
+
+
+def square_distance(src, dst):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src**2, -1).view(B, N, 1)
+    dist += torch.sum(dst**2, -1).view(B, 1, M)
+    return dist
+
+
+def index_points(points, idx):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = (
+        torch.arange(B, dtype=torch.long)
+        .to(device)
+        .view(view_shape)
+        .repeat(repeat_shape)
+    )
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+
+class KNearNeighbors(nn.Module):
+    """
+    Find the k nearest neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(KNearNeighbors, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, pos, centroids):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+        """
+        center_pos = index_points(pos, centroids)
+        sqrdists = square_distance(center_pos, pos)
+        group_idx = sqrdists.argsort(dim=-1)[:, :, : self.n_neighbor]
+        return group_idx
+
+
+class KNNGraphBuilder(nn.Module):
+    """
+    Build NN graph
+    """
+
+    def __init__(self, n_neighbor):
+        super(KNNGraphBuilder, self).__init__()
+        self.n_neighbor = n_neighbor
+        self.knn = KNearNeighbors(n_neighbor)
+
+    def forward(self, pos, centroids, feat=None):
+        dev = pos.device
+        group_idx = self.knn(pos, centroids)
+        B, N, _ = pos.shape
+        glist = []
+        for i in range(B):
+            center = torch.zeros((N)).to(dev)
+            center[centroids[i]] = 1
+            src = group_idx[i].contiguous().view(-1)
+            dst = (
+                centroids[i]
+                .view(-1, 1)
+                .repeat(
+                    1, min(self.n_neighbor, src.shape[0] // centroids.shape[1])
+                )
+                .view(-1)
+            )
+
+            unified = torch.cat([src, dst])
+            uniq, inv_idx = torch.unique(unified, return_inverse=True)
+            src_idx = inv_idx[: src.shape[0]]
+            dst_idx = inv_idx[src.shape[0] :]
+
+            g = dgl.graph((src_idx, dst_idx))
+            g.ndata["pos"] = pos[i][uniq]
+            g.ndata["center"] = center[uniq]
+            if feat is not None:
+                g.ndata["feat"] = feat[i][uniq]
+            glist.append(g)
+        bg = dgl.batch(glist)
+        return bg
+
+
+class RelativePositionMessage(nn.Module):
+    """
+    Compute the input feature from neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(RelativePositionMessage, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, edges):
+        pos = edges.src["pos"] - edges.dst["pos"]
+        if "feat" in edges.src:
+            res = torch.cat([pos, edges.src["feat"]], 1)
+        else:
+            res = pos
+        return {"agg_feat": res}
+
+
+class KNNConv(nn.Module):
+    """
+    Feature aggregation
+    """
+
+    def __init__(self, sizes, batch_size):
+        super(KNNConv, self).__init__()
+        self.batch_size = batch_size
+        self.conv = nn.ModuleList()
+        self.bn = nn.ModuleList()
+        for i in range(1, len(sizes)):
+            self.conv.append(nn.Conv2d(sizes[i - 1], sizes[i], 1))
+            self.bn.append(nn.BatchNorm2d(sizes[i]))
+
+    def forward(self, nodes):
+        shape = nodes.mailbox["agg_feat"].shape
+        h = (
+            nodes.mailbox["agg_feat"]
+            .view(self.batch_size, -1, shape[1], shape[2])
+            .permute(0, 3, 2, 1)
+        )
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h, 2)[0]
+        feat_dim = h.shape[1]
+        h = h.permute(0, 2, 1).reshape(-1, feat_dim)
+        return {"new_feat": h}
+
+    def group_all(self, pos, feat):
+        """
+        Feature aggregation and pooling for the non-sampling layer
+        """
+        if feat is not None:
+            h = torch.cat([pos, feat], 2)
+        else:
+            h = pos
+        B, N, D = h.shape
+        _, _, C = pos.shape
+        new_pos = torch.zeros(B, 1, C)
+        h = h.permute(0, 2, 1).view(B, -1, N, 1)
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h[:, :, :, 0], 2)[0]  # [B,D]
+        return new_pos, h
+
+
+class TransitionDown(nn.Module):
+    """
+    The Transition Down Module
+    """
+
+    def __init__(self, n_points, batch_size, mlp_sizes, n_neighbors=64):
+        super(TransitionDown, self).__init__()
+        self.n_points = n_points
+        self.frnn_graph = KNNGraphBuilder(n_neighbors)
+        self.message = RelativePositionMessage(n_neighbors)
+        self.conv = KNNConv(mlp_sizes, batch_size)
+        self.batch_size = batch_size
+
+    def forward(self, pos, feat):
+        centroids = farthest_point_sampler(pos, self.n_points)
+        g = self.frnn_graph(pos, centroids, feat)
+        g.update_all(self.message, self.conv)
+
+        mask = g.ndata["center"] == 1
+        pos_dim = g.ndata["pos"].shape[-1]
+        feat_dim = g.ndata["new_feat"].shape[-1]
+        pos_res = g.ndata["pos"][mask].view(self.batch_size, -1, pos_dim)
+        feat_res = g.ndata["new_feat"][mask].view(self.batch_size, -1, feat_dim)
+        return pos_res, feat_res
+
+
+class FeaturePropagation(nn.Module):
+    """
+    The FeaturePropagation Layer
+    """
+
+    def __init__(self, input_dims, sizes):
+        super(FeaturePropagation, self).__init__()
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+
+        sizes = [input_dims] + sizes
+        for i in range(1, len(sizes)):
+            self.convs.append(nn.Conv1d(sizes[i - 1], sizes[i], 1))
+            self.bns.append(nn.BatchNorm1d(sizes[i]))
+
+    def forward(self, x1, x2, feat1, feat2):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+            Input:
+                x1: input points position data, [B, N, C]
+                x2: sampled input points position data, [B, S, C]
+                feat1: input points data, [B, N, D]
+                feat2: input points data, [B, S, D]
+            Return:
+                new_feat: upsampled points data, [B, D', N]
+        """
+        B, N, C = x1.shape
+        _, S, _ = x2.shape
+
+        if S == 1:
+            interpolated_feat = feat2.repeat(1, N, 1)
+        else:
+            dists = square_distance(x1, x2)
+            dists, idx = dists.sort(dim=-1)
+            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]
+
+            dist_recip = 1.0 / (dists + 1e-8)
+            norm = torch.sum(dist_recip, dim=2, keepdim=True)
+            weight = dist_recip / norm
+            interpolated_feat = torch.sum(
+                index_points(feat2, idx) * weight.view(B, N, 3, 1), dim=2
+            )
+
+        if feat1 is not None:
+            new_feat = torch.cat([feat1, interpolated_feat], dim=-1)
+        else:
+            new_feat = interpolated_feat
+
+        new_feat = new_feat.permute(0, 2, 1)  # [B, D, S]
+        for i, conv in enumerate(self.convs):
+            bn = self.bns[i]
+            new_feat = F.relu(bn(conv(new_feat)))
+        return new_feat
+
+
+class SwapAxes(nn.Module):
+    def __init__(self, dim1=1, dim2=2):
+        super(SwapAxes, self).__init__()
+        self.dim1 = dim1
+        self.dim2 = dim2
+
+    def forward(self, x):
+        return x.transpose(self.dim1, self.dim2)
+
+
+class TransitionUp(nn.Module):
+    """
+    The Transition Up Module
+    """
+
+    def __init__(self, dim1, dim2, dim_out):
+        super(TransitionUp, self).__init__()
+        self.fc1 = nn.Sequential(
+            nn.Linear(dim1, dim_out),
+            SwapAxes(),
+            nn.BatchNorm1d(dim_out),  # TODO
+            SwapAxes(),
+            nn.ReLU(),
+        )
+        self.fc2 = nn.Sequential(
+            nn.Linear(dim2, dim_out),
+            SwapAxes(),
+            nn.BatchNorm1d(dim_out),  # TODO
+            SwapAxes(),
+            nn.ReLU(),
+        )
+        self.fp = FeaturePropagation(-1, [])
+
+    def forward(self, pos1, feat1, pos2, feat2):
+        h1 = self.fc1(feat1)
+        h2 = self.fc2(feat2)
+        h1 = self.fp(pos2, pos1, None, h1).transpose(1, 2)
+        return h1 + h2

examples/pytorch/pointcloud/point_transformer/point_transformer.py

@@ -0,0 +1,243 @@
+import numpy as np
+import torch
+from helper import index_points, square_distance, TransitionDown, TransitionUp
+from torch import nn
+
+"""
+Part of the code are adapted from
+https://github.com/qq456cvb/Point-Transformers
+"""
+
+
+class PointTransformerBlock(nn.Module):
+    def __init__(self, input_dim, n_neighbors, transformer_dim=None):
+        super(PointTransformerBlock, self).__init__()
+        if transformer_dim is None:
+            transformer_dim = input_dim
+        self.fc1 = nn.Linear(input_dim, transformer_dim)
+        self.fc2 = nn.Linear(transformer_dim, input_dim)
+        self.fc_delta = nn.Sequential(
+            nn.Linear(3, transformer_dim),
+            nn.ReLU(),
+            nn.Linear(transformer_dim, transformer_dim),
+        )
+        self.fc_gamma = nn.Sequential(
+            nn.Linear(transformer_dim, transformer_dim),
+            nn.ReLU(),
+            nn.Linear(transformer_dim, transformer_dim),
+        )
+        self.w_qs = nn.Linear(transformer_dim, transformer_dim, bias=False)
+        self.w_ks = nn.Linear(transformer_dim, transformer_dim, bias=False)
+        self.w_vs = nn.Linear(transformer_dim, transformer_dim, bias=False)
+        self.n_neighbors = n_neighbors
+
+    def forward(self, x, pos):
+        dists = square_distance(pos, pos)
+        knn_idx = dists.argsort()[:, :, : self.n_neighbors]  # b x n x k
+        knn_pos = index_points(pos, knn_idx)
+
+        h = self.fc1(x)
+        q, k, v = (
+            self.w_qs(h),
+            index_points(self.w_ks(h), knn_idx),
+            index_points(self.w_vs(h), knn_idx),
+        )
+
+        pos_enc = self.fc_delta(pos[:, :, None] - knn_pos)  # b x n x k x f
+
+        attn = self.fc_gamma(q[:, :, None] - k + pos_enc)
+        attn = torch.softmax(
+            attn / np.sqrt(k.size(-1)), dim=-2
+        )  # b x n x k x f
+
+        res = torch.einsum("bmnf,bmnf->bmf", attn, v + pos_enc)
+        res = self.fc2(res) + x
+        return res, attn
+
+
+class PointTransformer(nn.Module):
+    def __init__(
+        self,
+        n_points,
+        batch_size,
+        feature_dim=3,
+        n_blocks=4,
+        downsampling_rate=4,
+        hidden_dim=32,
+        transformer_dim=None,
+        n_neighbors=16,
+    ):
+        super(PointTransformer, self).__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.ptb = PointTransformerBlock(
+            hidden_dim, n_neighbors, transformer_dim
+        )
+        self.transition_downs = nn.ModuleList()
+        self.transformers = nn.ModuleList()
+        for i in range(n_blocks):
+            block_hidden_dim = hidden_dim * 2 ** (i + 1)
+            block_n_points = n_points // (downsampling_rate ** (i + 1))
+            self.transition_downs.append(
+                TransitionDown(
+                    block_n_points,
+                    batch_size,
+                    [
+                        block_hidden_dim // 2 + 3,
+                        block_hidden_dim,
+                        block_hidden_dim,
+                    ],
+                    n_neighbors=n_neighbors,
+                )
+            )
+            self.transformers.append(
+                PointTransformerBlock(
+                    block_hidden_dim, n_neighbors, transformer_dim
+                )
+            )
+
+    def forward(self, x):
+        if x.shape[-1] > 3:
+            pos = x[:, :, :3]
+        else:
+            pos = x
+
+        feat = x
+        h = self.fc(feat)
+        h, _ = self.ptb(h, pos)
+
+        hidden_state = [(pos, h)]
+        for td, tf in zip(self.transition_downs, self.transformers):
+            pos, h = td(pos, h)
+            h, _ = tf(h, pos)
+            hidden_state.append((pos, h))
+
+        return h, hidden_state
+
+
+class PointTransformerCLS(nn.Module):
+    def __init__(
+        self,
+        out_classes,
+        batch_size,
+        n_points=1024,
+        feature_dim=3,
+        n_blocks=4,
+        downsampling_rate=4,
+        hidden_dim=32,
+        transformer_dim=None,
+        n_neighbors=16,
+    ):
+        super(PointTransformerCLS, self).__init__()
+        self.backbone = PointTransformer(
+            n_points,
+            batch_size,
+            feature_dim,
+            n_blocks,
+            downsampling_rate,
+            hidden_dim,
+            transformer_dim,
+            n_neighbors,
+        )
+        self.out = self.fc2 = nn.Sequential(
+            nn.Linear(hidden_dim * 2 ** (n_blocks), 256),
+            nn.ReLU(),
+            nn.Linear(256, 64),
+            nn.ReLU(),
+            nn.Linear(64, out_classes),
+        )
+
+    def forward(self, x):
+        h, _ = self.backbone(x)
+        out = self.out(torch.mean(h, dim=1))
+        return out
+
+
+class PointTransformerSeg(nn.Module):
+    def __init__(
+        self,
+        out_classes,
+        batch_size,
+        n_points=2048,
+        feature_dim=3,
+        n_blocks=4,
+        downsampling_rate=4,
+        hidden_dim=32,
+        transformer_dim=None,
+        n_neighbors=16,
+    ):
+        super().__init__()
+        self.backbone = PointTransformer(
+            n_points,
+            batch_size,
+            feature_dim,
+            n_blocks,
+            downsampling_rate,
+            hidden_dim,
+            transformer_dim,
+            n_neighbors,
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(32 * 2**n_blocks, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 32 * 2**n_blocks),
+        )
+        self.ptb = PointTransformerBlock(
+            32 * 2**n_blocks, n_neighbors, transformer_dim
+        )
+
+        self.n_blocks = n_blocks
+        self.transition_ups = nn.ModuleList()
+        self.transformers = nn.ModuleList()
+        for i in reversed(range(n_blocks)):
+            block_hidden_dim = 32 * 2**i
+            self.transition_ups.append(
+                TransitionUp(
+                    block_hidden_dim * 2, block_hidden_dim, block_hidden_dim
+                )
+            )
+            self.transformers.append(
+                PointTransformerBlock(
+                    block_hidden_dim, n_neighbors, transformer_dim
+                )
+            )
+
+        self.out = nn.Sequential(
+            nn.Linear(32 + 16, 64),
+            nn.ReLU(),
+            nn.Linear(64, 64),
+            nn.ReLU(),
+            nn.Linear(64, out_classes),
+        )
+
+    def forward(self, x, cat_vec=None):
+        _, hidden_state = self.backbone(x)
+        pos, h = hidden_state[-1]
+        h, _ = self.ptb(self.fc(h), pos)
+
+        for i in range(self.n_blocks):
+            h = self.transition_ups[i](
+                pos, h, hidden_state[-i - 2][0], hidden_state[-i - 2][1]
+            )
+            pos = hidden_state[-i - 2][0]
+            h, _ = self.transformers[i](h, pos)
+        return self.out(torch.cat([h, cat_vec], dim=-1))
+
+
+class PartSegLoss(nn.Module):
+    def __init__(self, eps=0.2):
+        super(PartSegLoss, self).__init__()
+        self.eps = eps
+        self.loss = nn.CrossEntropyLoss()
+
+    def forward(self, logits, y):
+        num_classes = logits.shape[1]
+        logits = logits.permute(0, 2, 1).contiguous().view(-1, num_classes)
+        loss = self.loss(logits, y)
+        return loss

examples/pytorch/pointcloud/point_transformer/provider.py

@@ -0,0 +1,312 @@
+"""
+Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/master/provider.py
+"""
+import numpy as np
+
+
+def normalize_data(batch_data):
+    """Normalize the batch data, use coordinates of the block centered at origin,
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    B, N, C = batch_data.shape
+    normal_data = np.zeros((B, N, C))
+    for b in range(B):
+        pc = batch_data[b]
+        centroid = np.mean(pc, axis=0)
+        pc = pc - centroid
+        m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+        pc = pc / m
+        normal_data[b] = pc
+    return normal_data
+
+
+def shuffle_data(data, labels):
+    """Shuffle data and labels.
+    Input:
+      data: B,N,... numpy array
+      label: B,... numpy array
+    Return:
+      shuffled data, label and shuffle indices
+    """
+    idx = np.arange(len(labels))
+    np.random.shuffle(idx)
+    return data[idx, ...], labels[idx], idx
+
+
+def shuffle_points(batch_data):
+    """Shuffle orders of points in each point cloud -- changes FPS behavior.
+    Use the same shuffling idx for the entire batch.
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    idx = np.arange(batch_data.shape[1])
+    np.random.shuffle(idx)
+    return batch_data[:, idx, :]
+
+
+def rotate_point_cloud(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_z(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, sinval, 0], [-sinval, cosval, 0], [0, 0, 1]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_with_normal(batch_xyz_normal):
+    """Randomly rotate XYZ, normal point cloud.
+    Input:
+        batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
+    Output:
+        B,N,6, rotated XYZ, normal point cloud
+    """
+    for k in range(batch_xyz_normal.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_xyz_normal[k, :, 0:3]
+        shape_normal = batch_xyz_normal[k, :, 3:6]
+        batch_xyz_normal[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        batch_xyz_normal[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return batch_xyz_normal
+
+
+def rotate_perturbation_point_cloud_with_normal(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx6 array, original batch of point clouds and point normals
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
+        rotated_data[k, :, 3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx6 array, original batch of point clouds with normal
+      scalar, angle of rotation
+    Return:
+      BxNx6 array, rotated batch of point clouds iwth normal
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        rotated_data[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_perturbation_point_cloud(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
+    """Randomly jitter points. jittering is per point.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, jittered batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    assert clip > 0
+    jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1 * clip, clip)
+    jittered_data += batch_data
+    return jittered_data
+
+
+def shift_point_cloud(batch_data, shift_range=0.1):
+    """Randomly shift point cloud. Shift is per point cloud.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, shifted batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    shifts = np.random.uniform(-shift_range, shift_range, (B, 3))
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] += shifts[batch_index, :]
+    return batch_data
+
+
+def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25):
+    """Randomly scale the point cloud. Scale is per point cloud.
+    Input:
+        BxNx3 array, original batch of point clouds
+    Return:
+        BxNx3 array, scaled batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    scales = np.random.uniform(scale_low, scale_high, B)
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] *= scales[batch_index]
+    return batch_data
+
+
+def random_point_dropout(batch_pc, max_dropout_ratio=0.875):
+    """batch_pc: BxNx3"""
+    for b in range(batch_pc.shape[0]):
+        dropout_ratio = np.random.random() * max_dropout_ratio  # 0~0.875
+        drop_idx = np.where(
+            np.random.random((batch_pc.shape[1])) <= dropout_ratio
+        )[0]
+        if len(drop_idx) > 0:
+            dropout_ratio = (
+                np.random.random() * max_dropout_ratio
+            )  # 0~0.875 # not need
+            batch_pc[b, drop_idx, :] = batch_pc[
+                b, 0, :
+            ]  # set to the first point
+    return batch_pc

examples/pytorch/pointcloud/point_transformer/train_cls.py

@@ -0,0 +1,195 @@
+import argparse
+import os
+import time
+from functools import partial
+
+import provider
+import torch
+import torch.nn as nn
+import tqdm
+
+from dgl.data.utils import download, get_download_dir
+from ModelNetDataLoader import ModelNetDataLoader
+from point_transformer import PointTransformerCLS
+from torch.utils.data import DataLoader
+
+torch.backends.cudnn.enabled = False
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=200)
+parser.add_argument("--num-workers", type=int, default=8)
+parser.add_argument("--batch-size", type=int, default=16)
+parser.add_argument("--opt", type=str, default="adam")
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+data_filename = "modelnet40_normal_resampled.zip"
+download_path = os.path.join(get_download_dir(), data_filename)
+local_path = args.dataset_path or os.path.join(
+    get_download_dir(), "modelnet40_normal_resampled"
+)
+
+if not os.path.exists(local_path):
+    download(
+        "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip",
+        download_path,
+        verify_ssl=False,
+    )
+    from zipfile import ZipFile
+
+    with ZipFile(download_path) as z:
+        z.extractall(path=get_download_dir())
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    loss_f = nn.CrossEntropyLoss()
+    start_time = time.time()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label in tq:
+            data = data.data.numpy()
+            data = provider.random_point_dropout(data)
+            data[:, :, 0:3] = provider.random_scale_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.jitter_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.shift_point_cloud(data[:, :, 0:3])
+            data = torch.tensor(data)
+            label = label[:, 0]
+
+            num_examples = label.shape[0]
+            data, label = data.to(dev), label.to(dev).squeeze().long()
+            opt.zero_grad()
+            logits = net(data)
+            loss = loss_f(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            num_batches += 1
+            count += num_examples
+            loss = loss.item()
+            correct = (preds == label).sum().item()
+            total_loss += loss
+            total_correct += correct
+
+            tq.set_postfix(
+                {
+                    "AvgLoss": "%.5f" % (total_loss / num_batches),
+                    "AvgAcc": "%.5f" % (total_correct / count),
+                }
+            )
+    print(
+        "[Train] AvgLoss: {:.5}, AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_loss / num_batches,
+            total_correct / count,
+            time.time() - start_time,
+        )
+    )
+    scheduler.step()
+
+
+def evaluate(net, test_loader, dev):
+    net.eval()
+
+    total_correct = 0
+    count = 0
+    start_time = time.time()
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label in tq:
+                label = label[:, 0]
+                num_examples = label.shape[0]
+                data, label = data.to(dev), label.to(dev).squeeze().long()
+                logits = net(data)
+                _, preds = logits.max(1)
+
+                correct = (preds == label).sum().item()
+                total_correct += correct
+                count += num_examples
+
+                tq.set_postfix({"AvgAcc": "%.5f" % (total_correct / count)})
+    print(
+        "[Test]  AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_correct / count, time.time() - start_time
+        )
+    )
+    return total_correct / count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net = PointTransformerCLS(40, batch_size, feature_dim=6)
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+if args.opt == "sgd":
+    # The optimizer strategy described in paper:
+    opt = torch.optim.SGD(
+        net.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4
+    )
+    scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        opt, milestones=[120, 160], gamma=0.1
+    )
+elif args.opt == "adam":
+    # The optimizer strategy proposed by
+    # https://github.com/qq456cvb/Point-Transformers:
+    opt = torch.optim.Adam(
+        net.parameters(),
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-08,
+        weight_decay=1e-4,
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=50, gamma=0.3)
+
+train_dataset = ModelNetDataLoader(local_path, 1024, split="train")
+test_dataset = ModelNetDataLoader(local_path, 1024, split="test")
+train_loader = torch.utils.data.DataLoader(
+    train_dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    num_workers=num_workers,
+    drop_last=True,
+)
+test_loader = torch.utils.data.DataLoader(
+    test_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    drop_last=True,
+)
+
+best_test_acc = 0
+
+for epoch in range(args.num_epochs):
+    print("Epoch #{}: ".format(epoch))
+    train(net, opt, scheduler, train_loader, dev)
+    if (epoch + 1) % 1 == 0:
+        test_acc = evaluate(net, test_loader, dev)
+        if test_acc > best_test_acc:
+            best_test_acc = test_acc
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print("Current test acc: %.5f (best: %.5f)" % (test_acc, best_test_acc))
+    print()

examples/pytorch/pointcloud/point_transformer/train_partseg.py

@@ -0,0 +1,330 @@
+import argparse
+import time
+from functools import partial
+
+import dgl
+
+import numpy as np
+import torch
+import torch.optim as optim
+import tqdm
+from point_transformer import PartSegLoss, PointTransformerSeg
+from ShapeNet import ShapeNet
+from torch.utils.data import DataLoader
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=250)
+parser.add_argument("--num-workers", type=int, default=8)
+parser.add_argument("--batch-size", type=int, default=16)
+parser.add_argument("--tensorboard", action="store_true")
+parser.add_argument("--opt", type=str, default="adam")
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+
+def collate(samples):
+    graphs, cat = map(list, zip(*samples))
+    return dgl.batch(graphs), cat
+
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    start = time.time()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label, cat in tq:
+            num_examples = data.shape[0]
+            data = data.to(dev, dtype=torch.float)
+            label = label.to(dev, dtype=torch.long).view(-1)
+            opt.zero_grad()
+            cat_ind = [category_list.index(c) for c in cat]
+            # An one-hot encoding for the object category
+            cat_tensor = (
+                torch.tensor(eye_mat[cat_ind])
+                .to(dev, dtype=torch.float)
+                .repeat(1, 2048)
+            )
+            cat_tensor = cat_tensor.view(num_examples, -1, 16)
+            logits = net(data, cat_tensor).permute(0, 2, 1)
+            loss = L(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            count += num_examples * 2048
+            loss = loss.item()
+            total_loss += loss
+            num_batches += 1
+            correct = (preds.view(-1) == label).sum().item()
+            total_correct += correct
+
+            AvgLoss = total_loss / num_batches
+            AvgAcc = total_correct / count
+
+            tq.set_postfix(
+                {"AvgLoss": "%.5f" % AvgLoss, "AvgAcc": "%.5f" % AvgAcc}
+            )
+    scheduler.step()
+    end = time.time()
+    print(
+        "[Train] AvgLoss: {:.5}, AvgAcc: {:.5}, Time: {:.5}s".format(
+            total_loss / num_batches, total_correct / count, end - start
+        )
+    )
+    return data, preds, AvgLoss, AvgAcc, end - start
+
+
+def mIoU(preds, label, cat, cat_miou, seg_classes):
+    for i in range(preds.shape[0]):
+        shape_iou = 0
+        n = len(seg_classes[cat[i]])
+        for cls in seg_classes[cat[i]]:
+            pred_set = set(np.where(preds[i, :] == cls)[0])
+            label_set = set(np.where(label[i, :] == cls)[0])
+            union = len(pred_set.union(label_set))
+            inter = len(pred_set.intersection(label_set))
+            if union == 0:
+                shape_iou += 1
+            else:
+                shape_iou += inter / union
+        shape_iou /= n
+        cat_miou[cat[i]][0] += shape_iou
+        cat_miou[cat[i]][1] += 1
+
+    return cat_miou
+
+
+def evaluate(net, test_loader, dev, per_cat_verbose=False):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.eval()
+
+    cat_miou = {}
+    for k in shapenet.seg_classes.keys():
+        cat_miou[k] = [0, 0]
+    miou = 0
+    count = 0
+    per_cat_miou = 0
+    per_cat_count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label, cat in tq:
+                num_examples = data.shape[0]
+                data = data.to(dev, dtype=torch.float)
+                label = label.to(dev, dtype=torch.long)
+                cat_ind = [category_list.index(c) for c in cat]
+                cat_tensor = (
+                    torch.tensor(eye_mat[cat_ind])
+                    .to(dev, dtype=torch.float)
+                    .repeat(1, 2048)
+                )
+                cat_tensor = cat_tensor.view(num_examples, -1, 16)
+                logits = net(data, cat_tensor).permute(0, 2, 1)
+                _, preds = logits.max(1)
+
+                cat_miou = mIoU(
+                    preds.cpu().numpy(),
+                    label.view(num_examples, -1).cpu().numpy(),
+                    cat,
+                    cat_miou,
+                    shapenet.seg_classes,
+                )
+                for _, v in cat_miou.items():
+                    if v[1] > 0:
+                        miou += v[0]
+                        count += v[1]
+                        per_cat_miou += v[0] / v[1]
+                        per_cat_count += 1
+                tq.set_postfix(
+                    {
+                        "mIoU": "%.5f" % (miou / count),
+                        "per Category mIoU": "%.5f"
+                        % (per_cat_miou / per_cat_count),
+                    }
+                )
+    print(
+        "[Test] mIoU: %.5f, per Category mIoU: %.5f"
+        % (miou / count, per_cat_miou / per_cat_count)
+    )
+    if per_cat_verbose:
+        print("-" * 60)
+        print("Per-Category mIoU:")
+        for k, v in cat_miou.items():
+            if v[1] > 0:
+                print("%s mIoU=%.5f" % (k, v[0] / v[1]))
+            else:
+                print("%s mIoU=%.5f" % (k, 1))
+        print("-" * 60)
+    return miou / count, per_cat_miou / per_cat_count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net = PointTransformerSeg(50, batch_size)
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+if args.opt == "sgd":
+    # The optimizer strategy described in paper:
+    opt = torch.optim.SGD(
+        net.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4
+    )
+    scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        opt, milestones=[120, 160], gamma=0.1
+    )
+elif args.opt == "adam":
+    # The optimizer strategy proposed by
+    # https://github.com/qq456cvb/Point-Transformers:
+    opt = torch.optim.Adam(
+        net.parameters(),
+        lr=1e-3,
+        betas=(0.9, 0.999),
+        eps=1e-08,
+        weight_decay=1e-4,
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=50, gamma=0.3)
+
+L = PartSegLoss()
+
+shapenet = ShapeNet(2048, normal_channel=False)
+
+train_loader = CustomDataLoader(shapenet.trainval())
+test_loader = CustomDataLoader(shapenet.test())
+
+# Tensorboard
+if args.tensorboard:
+    import torchvision
+    from torch.utils.tensorboard import SummaryWriter
+    from torchvision import datasets, transforms
+
+    writer = SummaryWriter()
+# Select 50 distinct colors for different parts
+color_map = torch.tensor(
+    [
+        [47, 79, 79],
+        [139, 69, 19],
+        [112, 128, 144],
+        [85, 107, 47],
+        [139, 0, 0],
+        [128, 128, 0],
+        [72, 61, 139],
+        [0, 128, 0],
+        [188, 143, 143],
+        [60, 179, 113],
+        [205, 133, 63],
+        [0, 139, 139],
+        [70, 130, 180],
+        [205, 92, 92],
+        [154, 205, 50],
+        [0, 0, 139],
+        [50, 205, 50],
+        [250, 250, 250],
+        [218, 165, 32],
+        [139, 0, 139],
+        [10, 10, 10],
+        [176, 48, 96],
+        [72, 209, 204],
+        [153, 50, 204],
+        [255, 69, 0],
+        [255, 145, 0],
+        [0, 0, 205],
+        [255, 255, 0],
+        [0, 255, 0],
+        [233, 150, 122],
+        [220, 20, 60],
+        [0, 191, 255],
+        [160, 32, 240],
+        [192, 192, 192],
+        [173, 255, 47],
+        [218, 112, 214],
+        [216, 191, 216],
+        [255, 127, 80],
+        [255, 0, 255],
+        [100, 149, 237],
+        [128, 128, 128],
+        [221, 160, 221],
+        [144, 238, 144],
+        [123, 104, 238],
+        [255, 160, 122],
+        [175, 238, 238],
+        [238, 130, 238],
+        [127, 255, 212],
+        [255, 218, 185],
+        [255, 105, 180],
+    ]
+)
+# paint each point according to its pred
+
+
+def paint(batched_points):
+    B, N = batched_points.shape
+    colored = color_map[batched_points].squeeze(2)
+    return colored
+
+
+best_test_miou = 0
+best_test_per_cat_miou = 0
+
+for epoch in range(args.num_epochs):
+    print("Epoch #{}: ".format(epoch))
+    data, preds, AvgLoss, AvgAcc, training_time = train(
+        net, opt, scheduler, train_loader, dev
+    )
+    if (epoch + 1) % 5 == 0 or epoch == 0:
+        test_miou, test_per_cat_miou = evaluate(net, test_loader, dev, True)
+        if test_miou > best_test_miou:
+            best_test_miou = test_miou
+            best_test_per_cat_miou = test_per_cat_miou
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print(
+            "Current test mIoU: %.5f (best: %.5f), per-Category mIoU: %.5f (best: %.5f)"
+            % (
+                test_miou,
+                best_test_miou,
+                test_per_cat_miou,
+                best_test_per_cat_miou,
+            )
+        )
+    # Tensorboard
+    if args.tensorboard:
+        colored = paint(preds)
+        writer.add_mesh(
+            "data", vertices=data, colors=colored, global_step=epoch
+        )
+        writer.add_scalar(
+            "training time for one epoch", training_time, global_step=epoch
+        )
+        writer.add_scalar("AvgLoss", AvgLoss, global_step=epoch)
+        writer.add_scalar("AvgAcc", AvgAcc, global_step=epoch)
+        if (epoch + 1) % 5 == 0:
+            writer.add_scalar("test mIoU", test_miou, global_step=epoch)
+            writer.add_scalar(
+                "best test mIoU", best_test_miou, global_step=epoch
+            )
+    print()

examples/pytorch/pointcloud/pointnet/ModelNetDataLoader.py

@@ -0,0 +1,128 @@
+import os
+import warnings
+
+import numpy as np
+from torch.utils.data import Dataset
+
+warnings.filterwarnings("ignore")
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def farthest_point_sample(point, npoint):
+    """
+    Farthest point sampler works as follows:
+    1. Initialize the sample set S with a random point
+    2. Pick point P not in S, which maximizes the distance d(P, S)
+    3. Repeat step 2 until |S| = npoint
+
+    Input:
+        xyz: pointcloud data, [N, D]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [npoint, D]
+    """
+    N, D = point.shape
+    xyz = point[:, :3]
+    centroids = np.zeros((npoint,))
+    distance = np.ones((N,)) * 1e10
+    farthest = np.random.randint(0, N)
+    for i in range(npoint):
+        centroids[i] = farthest
+        centroid = xyz[farthest, :]
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+    point = point[centroids.astype(np.int32)]
+    return point
+
+
+class ModelNetDataLoader(Dataset):
+    def __init__(
+        self,
+        root,
+        npoint=1024,
+        split="train",
+        fps=False,
+        normal_channel=True,
+        cache_size=15000,
+    ):
+        """
+        Input:
+            root: the root path to the local data files
+            npoint: number of points from each cloud
+            split: which split of the data, 'train' or 'test'
+            fps: whether to sample points with farthest point sampler
+            normal_channel: whether to use additional channel
+            cache_size: the cache size of in-memory point clouds
+        """
+        self.root = root
+        self.npoints = npoint
+        self.fps = fps
+        self.catfile = os.path.join(self.root, "modelnet40_shape_names.txt")
+
+        self.cat = [line.rstrip() for line in open(self.catfile)]
+        self.classes = dict(zip(self.cat, range(len(self.cat))))
+        self.normal_channel = normal_channel
+
+        shape_ids = {}
+        shape_ids["train"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_train.txt"))
+        ]
+        shape_ids["test"] = [
+            line.rstrip()
+            for line in open(os.path.join(self.root, "modelnet40_test.txt"))
+        ]
+
+        assert split == "train" or split == "test"
+        shape_names = ["_".join(x.split("_")[0:-1]) for x in shape_ids[split]]
+        # list of (shape_name, shape_txt_file_path) tuple
+        self.datapath = [
+            (
+                shape_names[i],
+                os.path.join(self.root, shape_names[i], shape_ids[split][i])
+                + ".txt",
+            )
+            for i in range(len(shape_ids[split]))
+        ]
+        print("The size of %s data is %d" % (split, len(self.datapath)))
+
+        self.cache_size = cache_size
+        self.cache = {}
+
+    def __len__(self):
+        return len(self.datapath)
+
+    def _get_item(self, index):
+        if index in self.cache:
+            point_set, cls = self.cache[index]
+        else:
+            fn = self.datapath[index]
+            cls = self.classes[self.datapath[index][0]]
+            cls = np.array([cls]).astype(np.int32)
+            point_set = np.loadtxt(fn[1], delimiter=",").astype(np.float32)
+            if self.fps:
+                point_set = farthest_point_sample(point_set, self.npoints)
+            else:
+                point_set = point_set[0 : self.npoints, :]
+
+            point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
+
+            if not self.normal_channel:
+                point_set = point_set[:, 0:3]
+
+            if len(self.cache) < self.cache_size:
+                self.cache[index] = (point_set, cls)
+
+        return point_set, cls
+
+    def __getitem__(self, index):
+        return self._get_item(index)

examples/pytorch/pointcloud/pointnet/README.md

@@ -0,0 +1,49 @@
+PointNet and PointNet++ for Point Cloud Classification and Segmentation
+====
+
+This is a reproduction of the papers
+- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://arxiv.org/abs/1612.00593).
+- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://arxiv.org/abs/1706.02413).
+
+# Performance
+
+## Classification
+| Model           | Dataset    | Metric   | Score - PyTorch  | Score - DGL | Time(s) - PyTorch | Time(s) - DGL |
+|-----------------|------------|----------|------------------|-------------|-------------------|---------------|
+| PointNet        | ModelNet40 | Accuracy | 89.2(Official)   | 89.3        | 181.8             | 95.0          |
+| PointNet++(SSG) | ModelNet40 | Accuracy | 92.4             | 93.3        | 182.6             | 133.7         |
+| PointNet++(MSG) | ModelNet40 | Accuracy | 92.8             | 93.3        | 383.6             | 240.5         |
+
+## Part Segmentation
+
+| Model           | Dataset    | Metric   | Score - PyTorch | Score - DGL | Time(s) - PyTorch | Time(s) - DGL |
+|-----------------|------------|----------|-----------------|-------------|-------------------|---------------|
+| PointNet        | ShapeNet   | mIoU     | 84.3            | 83.6        | 251.6             | 234.0         |
+| PointNet++(SSG) | ShapeNet   | mIoU     | 84.9            | 84.5        | 361.7             | 240.1         |
+| PointNet++(MSG) | ShapeNet   | mIoU     | 85.4            | 84.6        | 817.3             | 821.8         |
+
++ Score - PyTorch are collected from [this repo](https://github.com/yanx27/Pointnet_Pointnet2_pytorch).
++ Time(s) are the average training time per epoch, measured on EC2 g4dn.4xlarge instance w/ Tesla T4 GPU.
+# How to Run
+
+For point cloud classification, run with
+
+```python
+python train_cls.py
+```
+
+For point cloud part-segmentation, run with
+
+```python
+python train_partseg.py
+```
+
+## To Visualize Part Segmentation in Tensorboard
+![Screenshot](vis.png)
+First ``pip install tensorboard``
+then run
+```python 
+python train_partseg.py --tensorboard
+```
+To display in Tensorboard, run 
+``tensorboard --logdir=runs``

examples/pytorch/pointcloud/pointnet/ShapeNet.py

@@ -0,0 +1,161 @@
+import json
+import os
+from zipfile import ZipFile
+
+import dgl
+
+import numpy as np
+import tqdm
+from dgl.data.utils import download, get_download_dir
+from scipy.sparse import csr_matrix
+from torch.utils.data import Dataset
+
+
+class ShapeNet(object):
+    def __init__(self, num_points=2048, normal_channel=True):
+        self.num_points = num_points
+        self.normal_channel = normal_channel
+
+        SHAPENET_DOWNLOAD_URL = "https://shapenet.cs.stanford.edu/media/shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        download_path = get_download_dir()
+        data_filename = (
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal.zip"
+        )
+        data_path = os.path.join(
+            download_path,
+            "shapenetcore_partanno_segmentation_benchmark_v0_normal",
+        )
+        if not os.path.exists(data_path):
+            local_path = os.path.join(download_path, data_filename)
+            if not os.path.exists(local_path):
+                download(SHAPENET_DOWNLOAD_URL, local_path, verify_ssl=False)
+            with ZipFile(local_path) as z:
+                z.extractall(path=download_path)
+
+        synset_file = "synsetoffset2category.txt"
+        with open(os.path.join(data_path, synset_file)) as f:
+            synset = [t.split("\n")[0].split("\t") for t in f.readlines()]
+        self.synset_dict = {}
+        for syn in synset:
+            self.synset_dict[syn[1]] = syn[0]
+        self.seg_classes = {
+            "Airplane": [0, 1, 2, 3],
+            "Bag": [4, 5],
+            "Cap": [6, 7],
+            "Car": [8, 9, 10, 11],
+            "Chair": [12, 13, 14, 15],
+            "Earphone": [16, 17, 18],
+            "Guitar": [19, 20, 21],
+            "Knife": [22, 23],
+            "Lamp": [24, 25, 26, 27],
+            "Laptop": [28, 29],
+            "Motorbike": [30, 31, 32, 33, 34, 35],
+            "Mug": [36, 37],
+            "Pistol": [38, 39, 40],
+            "Rocket": [41, 42, 43],
+            "Skateboard": [44, 45, 46],
+            "Table": [47, 48, 49],
+        }
+
+        train_split_json = "shuffled_train_file_list.json"
+        val_split_json = "shuffled_val_file_list.json"
+        test_split_json = "shuffled_test_file_list.json"
+        split_path = os.path.join(data_path, "train_test_split")
+        with open(os.path.join(split_path, train_split_json)) as f:
+            tmp = f.read()
+            self.train_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, val_split_json)) as f:
+            tmp = f.read()
+            self.val_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+        with open(os.path.join(split_path, test_split_json)) as f:
+            tmp = f.read()
+            self.test_file_list = [
+                os.path.join(data_path, t.replace("shape_data/", "") + ".txt")
+                for t in json.loads(tmp)
+            ]
+
+    def train(self):
+        return ShapeNetDataset(
+            self, "train", self.num_points, self.normal_channel
+        )
+
+    def valid(self):
+        return ShapeNetDataset(
+            self, "valid", self.num_points, self.normal_channel
+        )
+
+    def trainval(self):
+        return ShapeNetDataset(
+            self, "trainval", self.num_points, self.normal_channel
+        )
+
+    def test(self):
+        return ShapeNetDataset(
+            self, "test", self.num_points, self.normal_channel
+        )
+
+
+class ShapeNetDataset(Dataset):
+    def __init__(self, shapenet, mode, num_points, normal_channel=True):
+        super(ShapeNetDataset, self).__init__()
+        self.mode = mode
+        self.num_points = num_points
+        if not normal_channel:
+            self.dim = 3
+        else:
+            self.dim = 6
+
+        if mode == "train":
+            self.file_list = shapenet.train_file_list
+        elif mode == "valid":
+            self.file_list = shapenet.val_file_list
+        elif mode == "test":
+            self.file_list = shapenet.test_file_list
+        elif mode == "trainval":
+            self.file_list = shapenet.train_file_list + shapenet.val_file_list
+        else:
+            raise "Not supported `mode`"
+
+        data_list = []
+        label_list = []
+        category_list = []
+        print("Loading data from split " + self.mode)
+        for fn in tqdm.tqdm(self.file_list, ascii=True):
+            with open(fn) as f:
+                data = np.array(
+                    [t.split("\n")[0].split(" ") for t in f.readlines()]
+                ).astype(float)
+            data_list.append(data[:, 0 : self.dim])
+            label_list.append(data[:, 6].astype(int))
+            category_list.append(shapenet.synset_dict[fn.split("/")[-2]])
+        self.data = data_list
+        self.label = label_list
+        self.category = category_list
+
+    def translate(self, x, scale=(2 / 3, 3 / 2), shift=(-0.2, 0.2), size=3):
+        xyz1 = np.random.uniform(low=scale[0], high=scale[1], size=[size])
+        xyz2 = np.random.uniform(low=shift[0], high=shift[1], size=[size])
+        x = np.add(np.multiply(x, xyz1), xyz2).astype("float32")
+        return x
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        inds = np.random.choice(
+            self.data[i].shape[0], self.num_points, replace=True
+        )
+        x = self.data[i][inds, : self.dim]
+        y = self.label[i][inds]
+        cat = self.category[i]
+        if self.mode == "train":
+            x = self.translate(x, size=self.dim)
+        x = x.astype(float)
+        y = y.astype(int)
+        return x, y, cat

examples/pytorch/pointcloud/pointnet/pointnet2.py

@@ -0,0 +1,447 @@
+import dgl
+import dgl.function as fn
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.geometry import (
+    farthest_point_sampler,
+)  # dgl.geometry.pytorch -> dgl.geometry
+from torch.autograd import Variable
+
+"""
+Part of the code are adapted from
+https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+"""
+
+
+def square_distance(src, dst):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src**2, -1).view(B, N, 1)
+    dist += torch.sum(dst**2, -1).view(B, 1, M)
+    return dist
+
+
+def index_points(points, idx):
+    """
+    Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = (
+        torch.arange(B, dtype=torch.long)
+        .to(device)
+        .view(view_shape)
+        .repeat(repeat_shape)
+    )
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+
+class FixedRadiusNearNeighbors(nn.Module):
+    """
+    Ball Query - Find the neighbors with-in a fixed radius
+    """
+
+    def __init__(self, radius, n_neighbor):
+        super(FixedRadiusNearNeighbors, self).__init__()
+        self.radius = radius
+        self.n_neighbor = n_neighbor
+
+    def forward(self, pos, centroids):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+        """
+        device = pos.device
+        B, N, _ = pos.shape
+        center_pos = index_points(pos, centroids)
+        _, S, _ = center_pos.shape
+        group_idx = (
+            torch.arange(N, dtype=torch.long)
+            .to(device)
+            .view(1, 1, N)
+            .repeat([B, S, 1])
+        )
+        sqrdists = square_distance(center_pos, pos)
+        group_idx[sqrdists > self.radius**2] = N
+        group_idx = group_idx.sort(dim=-1)[0][:, :, : self.n_neighbor]
+        group_first = (
+            group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, self.n_neighbor])
+        )
+        mask = group_idx == N
+        group_idx[mask] = group_first[mask]
+        return group_idx
+
+
+class FixedRadiusNNGraph(nn.Module):
+    """
+    Build NN graph
+    """
+
+    def __init__(self, radius, n_neighbor):
+        super(FixedRadiusNNGraph, self).__init__()
+        self.radius = radius
+        self.n_neighbor = n_neighbor
+        self.frnn = FixedRadiusNearNeighbors(radius, n_neighbor)
+
+    def forward(self, pos, centroids, feat=None):
+        dev = pos.device
+        group_idx = self.frnn(pos, centroids)
+        B, N, _ = pos.shape
+        glist = []
+        for i in range(B):
+            center = torch.zeros((N)).to(dev)
+            center[centroids[i]] = 1
+            src = group_idx[i].contiguous().view(-1)
+            dst = centroids[i].view(-1, 1).repeat(1, self.n_neighbor).view(-1)
+
+            unified = torch.cat([src, dst])
+            uniq, inv_idx = torch.unique(unified, return_inverse=True)
+            src_idx = inv_idx[: src.shape[0]]
+            dst_idx = inv_idx[src.shape[0] :]
+
+            g = dgl.graph((src_idx, dst_idx))
+            g.ndata["pos"] = pos[i][uniq]
+            g.ndata["center"] = center[uniq]
+            if feat is not None:
+                g.ndata["feat"] = feat[i][uniq]
+            glist.append(g)
+        bg = dgl.batch(glist)
+        return bg
+
+
+class RelativePositionMessage(nn.Module):
+    """
+    Compute the input feature from neighbors
+    """
+
+    def __init__(self, n_neighbor):
+        super(RelativePositionMessage, self).__init__()
+        self.n_neighbor = n_neighbor
+
+    def forward(self, edges):
+        pos = edges.src["pos"] - edges.dst["pos"]
+        if "feat" in edges.src:
+            res = torch.cat([pos, edges.src["feat"]], 1)
+        else:
+            res = pos
+        return {"agg_feat": res}
+
+
+class PointNetConv(nn.Module):
+    """
+    Feature aggregation
+    """
+
+    def __init__(self, sizes, batch_size):
+        super(PointNetConv, self).__init__()
+        self.batch_size = batch_size
+        self.conv = nn.ModuleList()
+        self.bn = nn.ModuleList()
+        for i in range(1, len(sizes)):
+            self.conv.append(nn.Conv2d(sizes[i - 1], sizes[i], 1))
+            self.bn.append(nn.BatchNorm2d(sizes[i]))
+
+    def forward(self, nodes):
+        shape = nodes.mailbox["agg_feat"].shape
+        h = (
+            nodes.mailbox["agg_feat"]
+            .view(self.batch_size, -1, shape[1], shape[2])
+            .permute(0, 3, 2, 1)
+        )
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h, 2)[0]
+        feat_dim = h.shape[1]
+        h = h.permute(0, 2, 1).reshape(-1, feat_dim)
+        return {"new_feat": h}
+
+    def group_all(self, pos, feat):
+        """
+        Feature aggregation and pooling for the non-sampling layer
+        """
+        if feat is not None:
+            h = torch.cat([pos, feat], 2)
+        else:
+            h = pos
+        B, N, D = h.shape
+        _, _, C = pos.shape
+        new_pos = torch.zeros(B, 1, C)
+        h = h.permute(0, 2, 1).view(B, -1, N, 1)
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+        h = torch.max(h[:, :, :, 0], 2)[0]  # [B,D]
+        return new_pos, h
+
+
+class SAModule(nn.Module):
+    """
+    The Set Abstraction Layer
+    """
+
+    def __init__(
+        self,
+        npoints,
+        batch_size,
+        radius,
+        mlp_sizes,
+        n_neighbor=64,
+        group_all=False,
+    ):
+        super(SAModule, self).__init__()
+        self.group_all = group_all
+        if not group_all:
+            self.npoints = npoints
+            self.frnn_graph = FixedRadiusNNGraph(radius, n_neighbor)
+        self.message = RelativePositionMessage(n_neighbor)
+        self.conv = PointNetConv(mlp_sizes, batch_size)
+        self.batch_size = batch_size
+
+    def forward(self, pos, feat):
+        if self.group_all:
+            return self.conv.group_all(pos, feat)
+
+        centroids = farthest_point_sampler(pos, self.npoints)
+        g = self.frnn_graph(pos, centroids, feat)
+        g.update_all(self.message, self.conv)
+
+        mask = g.ndata["center"] == 1
+        pos_dim = g.ndata["pos"].shape[-1]
+        feat_dim = g.ndata["new_feat"].shape[-1]
+        pos_res = g.ndata["pos"][mask].view(self.batch_size, -1, pos_dim)
+        feat_res = g.ndata["new_feat"][mask].view(self.batch_size, -1, feat_dim)
+        return pos_res, feat_res
+
+
+class SAMSGModule(nn.Module):
+    """
+    The Set Abstraction Multi-Scale grouping Layer
+    """
+
+    def __init__(
+        self, npoints, batch_size, radius_list, n_neighbor_list, mlp_sizes_list
+    ):
+        super(SAMSGModule, self).__init__()
+        self.batch_size = batch_size
+        self.group_size = len(radius_list)
+
+        self.npoints = npoints
+        self.frnn_graph_list = nn.ModuleList()
+        self.message_list = nn.ModuleList()
+        self.conv_list = nn.ModuleList()
+        for i in range(self.group_size):
+            self.frnn_graph_list.append(
+                FixedRadiusNNGraph(radius_list[i], n_neighbor_list[i])
+            )
+            self.message_list.append(
+                RelativePositionMessage(n_neighbor_list[i])
+            )
+            self.conv_list.append(PointNetConv(mlp_sizes_list[i], batch_size))
+
+    def forward(self, pos, feat):
+        centroids = farthest_point_sampler(pos, self.npoints)
+        feat_res_list = []
+
+        for i in range(self.group_size):
+            g = self.frnn_graph_list[i](pos, centroids, feat)
+            g.update_all(self.message_list[i], self.conv_list[i])
+            mask = g.ndata["center"] == 1
+            pos_dim = g.ndata["pos"].shape[-1]
+            feat_dim = g.ndata["new_feat"].shape[-1]
+            if i == 0:
+                pos_res = g.ndata["pos"][mask].view(
+                    self.batch_size, -1, pos_dim
+                )
+            feat_res = g.ndata["new_feat"][mask].view(
+                self.batch_size, -1, feat_dim
+            )
+            feat_res_list.append(feat_res)
+
+        feat_res = torch.cat(feat_res_list, 2)
+        return pos_res, feat_res
+
+
+class PointNet2FP(nn.Module):
+    """
+    The Feature Propagation Layer
+    """
+
+    def __init__(self, input_dims, sizes):
+        super(PointNet2FP, self).__init__()
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+
+        sizes = [input_dims] + sizes
+        for i in range(1, len(sizes)):
+            self.convs.append(nn.Conv1d(sizes[i - 1], sizes[i], 1))
+            self.bns.append(nn.BatchNorm1d(sizes[i]))
+
+    def forward(self, x1, x2, feat1, feat2):
+        """
+        Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch
+            Input:
+                x1: input points position data, [B, N, C]
+                x2: sampled input points position data, [B, S, C]
+                feat1: input points data, [B, N, D]
+                feat2: input points data, [B, S, D]
+            Return:
+                new_feat: upsampled points data, [B, D', N]
+        """
+        B, N, C = x1.shape
+        _, S, _ = x2.shape
+
+        if S == 1:
+            interpolated_feat = feat2.repeat(1, N, 1)
+        else:
+            dists = square_distance(x1, x2)
+            dists, idx = dists.sort(dim=-1)
+            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]
+
+            dist_recip = 1.0 / (dists + 1e-8)
+            norm = torch.sum(dist_recip, dim=2, keepdim=True)
+            weight = dist_recip / norm
+            interpolated_feat = torch.sum(
+                index_points(feat2, idx) * weight.view(B, N, 3, 1), dim=2
+            )
+
+        if feat1 is not None:
+            new_feat = torch.cat([feat1, interpolated_feat], dim=-1)
+        else:
+            new_feat = interpolated_feat
+
+        new_feat = new_feat.permute(0, 2, 1)  # [B, D, S]
+        for i, conv in enumerate(self.convs):
+            bn = self.bns[i]
+            new_feat = F.relu(bn(conv(new_feat)))
+        return new_feat
+
+
+class PointNet2SSGCls(nn.Module):
+    def __init__(
+        self, output_classes, batch_size, input_dims=3, dropout_prob=0.4
+    ):
+        super(PointNet2SSGCls, self).__init__()
+        self.input_dims = input_dims
+
+        self.sa_module1 = SAModule(
+            512, batch_size, 0.2, [input_dims, 64, 64, 128]
+        )
+        self.sa_module2 = SAModule(
+            128, batch_size, 0.4, [128 + 3, 128, 128, 256]
+        )
+        self.sa_module3 = SAModule(
+            None, batch_size, None, [256 + 3, 256, 512, 1024], group_all=True
+        )
+
+        self.mlp1 = nn.Linear(1024, 512)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.drop1 = nn.Dropout(dropout_prob)
+
+        self.mlp2 = nn.Linear(512, 256)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.drop2 = nn.Dropout(dropout_prob)
+
+        self.mlp_out = nn.Linear(256, output_classes)
+
+    def forward(self, x):
+        if x.shape[-1] > 3:
+            pos = x[:, :, :3]
+            feat = x[:, :, 3:]
+        else:
+            pos = x
+            feat = None
+        pos, feat = self.sa_module1(pos, feat)
+        pos, feat = self.sa_module2(pos, feat)
+        _, h = self.sa_module3(pos, feat)
+
+        h = self.mlp1(h)
+        h = self.bn1(h)
+        h = F.relu(h)
+        h = self.drop1(h)
+        h = self.mlp2(h)
+        h = self.bn2(h)
+        h = F.relu(h)
+        h = self.drop2(h)
+
+        out = self.mlp_out(h)
+        return out
+
+
+class PointNet2MSGCls(nn.Module):
+    def __init__(
+        self, output_classes, batch_size, input_dims=3, dropout_prob=0.4
+    ):
+        super(PointNet2MSGCls, self).__init__()
+        self.input_dims = input_dims
+
+        self.sa_msg_module1 = SAMSGModule(
+            512,
+            batch_size,
+            [0.1, 0.2, 0.4],
+            [16, 32, 128],
+            [
+                [input_dims, 32, 32, 64],
+                [input_dims, 64, 64, 128],
+                [input_dims, 64, 96, 128],
+            ],
+        )
+        self.sa_msg_module2 = SAMSGModule(
+            128,
+            batch_size,
+            [0.2, 0.4, 0.8],
+            [32, 64, 128],
+            [
+                [320 + 3, 64, 64, 128],
+                [320 + 3, 128, 128, 256],
+                [320 + 3, 128, 128, 256],
+            ],
+        )
+        self.sa_module3 = SAModule(
+            None, batch_size, None, [640 + 3, 256, 512, 1024], group_all=True
+        )
+
+        self.mlp1 = nn.Linear(1024, 512)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.drop1 = nn.Dropout(dropout_prob)
+
+        self.mlp2 = nn.Linear(512, 256)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.drop2 = nn.Dropout(dropout_prob)
+
+        self.mlp_out = nn.Linear(256, output_classes)
+
+    def forward(self, x):
+        if x.shape[-1] > 3:
+            pos = x[:, :, :3]
+            feat = x[:, :, 3:]
+        else:
+            pos = x
+            feat = None
+        pos, feat = self.sa_msg_module1(pos, feat)
+        pos, feat = self.sa_msg_module2(pos, feat)
+        _, h = self.sa_module3(pos, feat)
+
+        h = self.mlp1(h)
+        h = self.bn1(h)
+        h = F.relu(h)
+        h = self.drop1(h)
+        h = self.mlp2(h)
+        h = self.bn2(h)
+        h = F.relu(h)
+        h = self.drop2(h)
+
+        out = self.mlp_out(h)
+        return out

examples/pytorch/pointcloud/pointnet/pointnet2_partseg.py

@@ -0,0 +1,119 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pointnet2 import PointNet2FP, SAModule, SAMSGModule
+from torch.autograd import Variable
+
+
+class PointNet2SSGPartSeg(nn.Module):
+    def __init__(self, output_classes, batch_size, input_dims=6):
+        super(PointNet2SSGPartSeg, self).__init__()
+        # if normal_channel == true, input_dims = 6+3
+        self.input_dims = input_dims
+
+        self.sa_module1 = SAModule(
+            512, batch_size, 0.2, [input_dims, 64, 64, 128], n_neighbor=32
+        )
+        self.sa_module2 = SAModule(
+            128, batch_size, 0.4, [128 + 3, 128, 128, 256]
+        )
+        self.sa_module3 = SAModule(
+            None, batch_size, None, [256 + 3, 256, 512, 1024], group_all=True
+        )
+
+        self.fp3 = PointNet2FP(1280, [256, 256])
+        self.fp2 = PointNet2FP(384, [256, 128])
+        # if normal_channel == true, 128+16+6+3
+        self.fp1 = PointNet2FP(128 + 16 + 6, [128, 128, 128])
+
+        self.conv1 = nn.Conv1d(128, 128, 1)
+        self.bn1 = nn.BatchNorm1d(128)
+        self.drop1 = nn.Dropout(0.5)
+        self.conv2 = nn.Conv1d(128, output_classes, 1)
+
+    def forward(self, x, cat_vec=None):
+        if x.shape[-1] > 3:
+            l0_pos = x[:, :, :3]
+            l0_feat = x
+        else:
+            l0_pos = x
+            l0_feat = x
+        # Set Abstraction layers
+        l1_pos, l1_feat = self.sa_module1(l0_pos, l0_feat)  # l1_feat: [B, N, D]
+        l2_pos, l2_feat = self.sa_module2(l1_pos, l1_feat)
+        l3_pos, l3_feat = self.sa_module3(l2_pos, l2_feat)  # [B, N, C], [B, D]
+        # Feature Propagation layers
+        l2_feat = self.fp3(
+            l2_pos, l3_pos, l2_feat, l3_feat.unsqueeze(1)
+        )  # l2_feat: [B, D, N]
+        l1_feat = self.fp2(l1_pos, l2_pos, l1_feat, l2_feat.permute(0, 2, 1))
+        l0_feat = torch.cat([cat_vec.permute(0, 2, 1), l0_pos, l0_feat], 2)
+        l0_feat = self.fp1(l0_pos, l1_pos, l0_feat, l1_feat.permute(0, 2, 1))
+        # FC layers
+        feat = F.relu(self.bn1(self.conv1(l0_feat)))
+        out = self.drop1(feat)
+        out = self.conv2(out)  # [B, output_classes, N]
+        return out
+
+
+class PointNet2MSGPartSeg(nn.Module):
+    def __init__(self, output_classes, batch_size, input_dims=6):
+        super(PointNet2MSGPartSeg, self).__init__()
+
+        self.sa_msg_module1 = SAMSGModule(
+            512,
+            batch_size,
+            [0.1, 0.2, 0.4],
+            [32, 64, 128],
+            [
+                [input_dims, 32, 32, 64],
+                [input_dims, 64, 64, 128],
+                [input_dims, 64, 96, 128],
+            ],
+        )
+        self.sa_msg_module2 = SAMSGModule(
+            128,
+            batch_size,
+            [0.4, 0.8],
+            [64, 128],
+            [
+                [128 + 128 + 64 + 3, 128, 128, 256],
+                [128 + 128 + 64 + 3, 128, 196, 256],
+            ],
+        )
+        self.sa_module3 = SAModule(
+            None, batch_size, None, [512 + 3, 256, 512, 1024], group_all=True
+        )
+
+        self.fp3 = PointNet2FP(1536, [256, 256])
+        self.fp2 = PointNet2FP(576, [256, 128])
+        # if normal_channel == true, 150 + 3
+        self.fp1 = PointNet2FP(150, [128, 128])
+
+        self.conv1 = nn.Conv1d(128, 128, 1)
+        self.bn1 = nn.BatchNorm1d(128)
+        self.drop1 = nn.Dropout(0.5)
+        self.conv2 = nn.Conv1d(128, output_classes, 1)
+
+    def forward(self, x, cat_vec=None):
+        if x.shape[-1] > 3:
+            l0_pos = x[:, :, :3]
+            l0_feat = x
+        else:
+            l0_pos = x
+            l0_feat = x
+        # Set Abstraction layers
+        l1_pos, l1_feat = self.sa_msg_module1(l0_pos, l0_feat)
+        l2_pos, l2_feat = self.sa_msg_module2(l1_pos, l1_feat)
+        l3_pos, l3_feat = self.sa_module3(l2_pos, l2_feat)
+        # Feature Propagation layers
+        l2_feat = self.fp3(l2_pos, l3_pos, l2_feat, l3_feat.unsqueeze(1))
+        l1_feat = self.fp2(l1_pos, l2_pos, l1_feat, l2_feat.permute(0, 2, 1))
+        l0_feat = torch.cat([cat_vec.permute(0, 2, 1), l0_pos, l0_feat], 2)
+        l0_feat = self.fp1(l0_pos, l1_pos, l0_feat, l1_feat.permute(0, 2, 1))
+        # FC layers
+        feat = F.relu(self.bn1(self.conv1(l0_feat)))
+        out = self.drop1(feat)
+        out = self.conv2(out)
+        return out

examples/pytorch/pointcloud/pointnet/pointnet_cls.py

@@ -0,0 +1,150 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+class PointNetCls(nn.Module):
+    def __init__(
+        self,
+        output_classes,
+        input_dims=3,
+        conv1_dim=64,
+        dropout_prob=0.5,
+        use_transform=True,
+    ):
+        super(PointNetCls, self).__init__()
+        self.input_dims = input_dims
+        self.conv1 = nn.ModuleList()
+        self.conv1.append(nn.Conv1d(input_dims, conv1_dim, 1))
+        self.conv1.append(nn.Conv1d(conv1_dim, conv1_dim, 1))
+        self.conv1.append(nn.Conv1d(conv1_dim, conv1_dim, 1))
+
+        self.bn1 = nn.ModuleList()
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+        self.bn1.append(nn.BatchNorm1d(conv1_dim))
+
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(nn.Conv1d(conv1_dim, conv1_dim * 2, 1))
+        self.conv2.append(nn.Conv1d(conv1_dim * 2, conv1_dim * 16, 1))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 2))
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 16))
+
+        self.maxpool = nn.MaxPool1d(conv1_dim * 16)
+        self.pool_feat_len = conv1_dim * 16
+
+        self.mlp3 = nn.ModuleList()
+        self.mlp3.append(nn.Linear(conv1_dim * 16, conv1_dim * 8))
+        self.mlp3.append(nn.Linear(conv1_dim * 8, conv1_dim * 4))
+
+        self.bn3 = nn.ModuleList()
+        self.bn3.append(nn.BatchNorm1d(conv1_dim * 8))
+        self.bn3.append(nn.BatchNorm1d(conv1_dim * 4))
+
+        self.dropout = nn.Dropout(0.3)
+        self.mlp_out = nn.Linear(conv1_dim * 4, output_classes)
+
+        self.use_transform = use_transform
+        if use_transform:
+            self.transform1 = TransformNet(input_dims)
+            self.trans_bn1 = nn.BatchNorm1d(input_dims)
+            self.transform2 = TransformNet(conv1_dim)
+            self.trans_bn2 = nn.BatchNorm1d(conv1_dim)
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        h = x.permute(0, 2, 1)
+        if self.use_transform:
+            trans = self.transform1(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn1(h))
+
+        for conv, bn in zip(self.conv1, self.bn1):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        if self.use_transform:
+            trans = self.transform2(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn2(h))
+
+        for conv, bn in zip(self.conv2, self.bn2):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.maxpool(h).view(-1, self.pool_feat_len)
+        for mlp, bn in zip(self.mlp3, self.bn3):
+            h = mlp(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.dropout(h)
+        out = self.mlp_out(h)
+        return out
+
+
+class TransformNet(nn.Module):
+    def __init__(self, input_dims=3, conv1_dim=64):
+        super(TransformNet, self).__init__()
+        self.conv = nn.ModuleList()
+        self.conv.append(nn.Conv1d(input_dims, conv1_dim, 1))
+        self.conv.append(nn.Conv1d(conv1_dim, conv1_dim * 2, 1))
+        self.conv.append(nn.Conv1d(conv1_dim * 2, conv1_dim * 16, 1))
+
+        self.bn = nn.ModuleList()
+        self.bn.append(nn.BatchNorm1d(conv1_dim))
+        self.bn.append(nn.BatchNorm1d(conv1_dim * 2))
+        self.bn.append(nn.BatchNorm1d(conv1_dim * 16))
+
+        self.maxpool = nn.MaxPool1d(conv1_dim * 16)
+        self.pool_feat_len = conv1_dim * 16
+
+        self.mlp2 = nn.ModuleList()
+        self.mlp2.append(nn.Linear(conv1_dim * 16, conv1_dim * 8))
+        self.mlp2.append(nn.Linear(conv1_dim * 8, conv1_dim * 4))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 8))
+        self.bn2.append(nn.BatchNorm1d(conv1_dim * 4))
+
+        self.input_dims = input_dims
+        self.mlp_out = nn.Linear(conv1_dim * 4, input_dims * input_dims)
+
+    def forward(self, h):
+        batch_size = h.shape[0]
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.maxpool(h).view(-1, self.pool_feat_len)
+        for mlp, bn in zip(self.mlp2, self.bn2):
+            h = mlp(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        out = self.mlp_out(h)
+
+        iden = Variable(
+            torch.from_numpy(
+                np.eye(self.input_dims).flatten().astype(np.float32)
+            )
+        )
+        iden = iden.view(1, self.input_dims * self.input_dims).repeat(
+            batch_size, 1
+        )
+        if out.is_cuda:
+            iden = iden.cuda()
+        out = out + iden
+        out = out.view(-1, self.input_dims, self.input_dims)
+        return out

examples/pytorch/pointcloud/pointnet/pointnet_partseg.py

@@ -0,0 +1,171 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+class PointNetPartSeg(nn.Module):
+    def __init__(
+        self, output_classes, input_dims=3, num_points=2048, use_transform=True
+    ):
+        super(PointNetPartSeg, self).__init__()
+        self.input_dims = input_dims
+
+        self.conv1 = nn.ModuleList()
+        self.conv1.append(nn.Conv1d(input_dims, 64, 1))
+        self.conv1.append(nn.Conv1d(64, 128, 1))
+        self.conv1.append(nn.Conv1d(128, 128, 1))
+
+        self.bn1 = nn.ModuleList()
+        self.bn1.append(nn.BatchNorm1d(64))
+        self.bn1.append(nn.BatchNorm1d(128))
+        self.bn1.append(nn.BatchNorm1d(128))
+
+        self.conv2 = nn.ModuleList()
+        self.conv2.append(nn.Conv1d(128, 512, 1))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(512))
+
+        self.conv_max = nn.Conv1d(512, 2048, 1)
+        self.bn_max = nn.BatchNorm1d(2048)
+
+        self.maxpool = nn.MaxPool1d(num_points)
+        self.pool_feat_len = 2048
+
+        self.conv3 = nn.ModuleList()
+        self.conv3.append(nn.Conv1d(2048 + 64 + 128 * 3 + 512 + 16, 256, 1))
+        self.conv3.append(nn.Conv1d(256, 256, 1))
+        self.conv3.append(nn.Conv1d(256, 128, 1))
+
+        self.bn3 = nn.ModuleList()
+        self.bn3.append(nn.BatchNorm1d(256))
+        self.bn3.append(nn.BatchNorm1d(256))
+        self.bn3.append(nn.BatchNorm1d(128))
+
+        self.conv_out = nn.Conv1d(128, output_classes, 1)
+
+        self.use_transform = use_transform
+        if use_transform:
+            self.transform1 = TransformNet(self.input_dims)
+            self.trans_bn1 = nn.BatchNorm1d(self.input_dims)
+            self.transform2 = TransformNet(128)
+            self.trans_bn2 = nn.BatchNorm1d(128)
+
+    def forward(self, x, cat_vec=None):
+        batch_size = x.shape[0]
+        h = x.permute(0, 2, 1)
+        num_points = h.shape[2]
+        if self.use_transform:
+            trans = self.transform1(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn1(h))
+
+        mid_feat = []
+        for conv, bn in zip(self.conv1, self.bn1):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+            mid_feat.append(h)
+
+        if self.use_transform:
+            trans = self.transform2(h)
+            h = h.transpose(2, 1)
+            h = torch.bmm(h, trans)
+            h = h.transpose(2, 1)
+            h = F.relu(self.trans_bn2(h))
+            mid_feat.append(h)
+
+        for conv, bn in zip(self.conv2, self.bn2):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+            mid_feat.append(h)
+
+        h = self.conv_max(h)
+        h = self.bn_max(h)
+        h = self.maxpool(h).view(batch_size, -1, 1).repeat(1, 1, num_points)
+        mid_feat.append(h)
+        if cat_vec is not None:
+            mid_feat.append(cat_vec)
+        h = torch.cat(mid_feat, 1)
+        for conv, bn in zip(self.conv3, self.bn3):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        out = self.conv_out(h)
+        return out
+
+
+class TransformNet(nn.Module):
+    def __init__(self, input_dims=3, num_points=2048):
+        super(TransformNet, self).__init__()
+        self.conv = nn.ModuleList()
+        self.conv.append(nn.Conv1d(input_dims, 64, 1))
+        self.conv.append(nn.Conv1d(64, 128, 1))
+        self.conv.append(nn.Conv1d(128, 1024, 1))
+
+        self.bn = nn.ModuleList()
+        self.bn.append(nn.BatchNorm1d(64))
+        self.bn.append(nn.BatchNorm1d(128))
+        self.bn.append(nn.BatchNorm1d(1024))
+
+        self.maxpool = nn.MaxPool1d(num_points)
+        self.pool_feat_len = 1024
+
+        self.mlp2 = nn.ModuleList()
+        self.mlp2.append(nn.Linear(1024, 512))
+        self.mlp2.append(nn.Linear(512, 256))
+
+        self.bn2 = nn.ModuleList()
+        self.bn2.append(nn.BatchNorm1d(512))
+        self.bn2.append(nn.BatchNorm1d(256))
+
+        self.input_dims = input_dims
+        self.mlp_out = nn.Linear(256, input_dims * input_dims)
+
+    def forward(self, h):
+        batch_size = h.shape[0]
+        for conv, bn in zip(self.conv, self.bn):
+            h = conv(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        h = self.maxpool(h).view(-1, self.pool_feat_len)
+        for mlp, bn in zip(self.mlp2, self.bn2):
+            h = mlp(h)
+            h = bn(h)
+            h = F.relu(h)
+
+        out = self.mlp_out(h)
+
+        iden = Variable(
+            torch.from_numpy(
+                np.eye(self.input_dims).flatten().astype(np.float32)
+            )
+        )
+        iden = iden.view(1, self.input_dims * self.input_dims).repeat(
+            batch_size, 1
+        )
+        if out.is_cuda:
+            iden = iden.cuda()
+        out = out + iden
+        out = out.view(-1, self.input_dims, self.input_dims)
+        return out
+
+
+class PartSegLoss(nn.Module):
+    def __init__(self, eps=0.2):
+        super(PartSegLoss, self).__init__()
+        self.eps = eps
+        self.loss = nn.CrossEntropyLoss()
+
+    def forward(self, logits, y):
+        num_classes = logits.shape[1]
+        logits = logits.permute(0, 2, 1).contiguous().view(-1, num_classes)
+        loss = self.loss(logits, y)
+        return loss

examples/pytorch/pointcloud/pointnet/provider.py

@@ -0,0 +1,312 @@
+"""
+Adapted from https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/master/provider.py
+"""
+import numpy as np
+
+
+def normalize_data(batch_data):
+    """Normalize the batch data, use coordinates of the block centered at origin,
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    B, N, C = batch_data.shape
+    normal_data = np.zeros((B, N, C))
+    for b in range(B):
+        pc = batch_data[b]
+        centroid = np.mean(pc, axis=0)
+        pc = pc - centroid
+        m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+        pc = pc / m
+        normal_data[b] = pc
+    return normal_data
+
+
+def shuffle_data(data, labels):
+    """Shuffle data and labels.
+    Input:
+      data: B,N,... numpy array
+      label: B,... numpy array
+    Return:
+      shuffled data, label and shuffle indices
+    """
+    idx = np.arange(len(labels))
+    np.random.shuffle(idx)
+    return data[idx, ...], labels[idx], idx
+
+
+def shuffle_points(batch_data):
+    """Shuffle orders of points in each point cloud -- changes FPS behavior.
+    Use the same shuffling idx for the entire batch.
+    Input:
+        BxNxC array
+    Output:
+        BxNxC array
+    """
+    idx = np.arange(batch_data.shape[1])
+    np.random.shuffle(idx)
+    return batch_data[:, idx, :]
+
+
+def rotate_point_cloud(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_z(batch_data):
+    """Randomly rotate the point clouds to augument the dataset
+    rotation is per shape based along up direction
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, sinval, 0], [-sinval, cosval, 0], [0, 0, 1]]
+        )
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_with_normal(batch_xyz_normal):
+    """Randomly rotate XYZ, normal point cloud.
+    Input:
+        batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
+    Output:
+        B,N,6, rotated XYZ, normal point cloud
+    """
+    for k in range(batch_xyz_normal.shape[0]):
+        rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_xyz_normal[k, :, 0:3]
+        shape_normal = batch_xyz_normal[k, :, 3:6]
+        batch_xyz_normal[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        batch_xyz_normal[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return batch_xyz_normal
+
+
+def rotate_perturbation_point_cloud_with_normal(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx6 array, original batch of point clouds and point normals
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
+        rotated_data[k, :, 3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
+    """Rotate the point cloud along up direction with certain angle.
+    Input:
+      BxNx6 array, original batch of point clouds with normal
+      scalar, angle of rotation
+    Return:
+      BxNx6 array, rotated batch of point clouds iwth normal
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        # rotation_angle = np.random.uniform() * 2 * np.pi
+        cosval = np.cos(rotation_angle)
+        sinval = np.sin(rotation_angle)
+        rotation_matrix = np.array(
+            [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]]
+        )
+        shape_pc = batch_data[k, :, 0:3]
+        shape_normal = batch_data[k, :, 3:6]
+        rotated_data[k, :, 0:3] = np.dot(
+            shape_pc.reshape((-1, 3)), rotation_matrix
+        )
+        rotated_data[k, :, 3:6] = np.dot(
+            shape_normal.reshape((-1, 3)), rotation_matrix
+        )
+    return rotated_data
+
+
+def rotate_perturbation_point_cloud(
+    batch_data, angle_sigma=0.06, angle_clip=0.18
+):
+    """Randomly perturb the point clouds by small rotations
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, rotated batch of point clouds
+    """
+    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
+    for k in range(batch_data.shape[0]):
+        angles = np.clip(
+            angle_sigma * np.random.randn(3), -angle_clip, angle_clip
+        )
+        Rx = np.array(
+            [
+                [1, 0, 0],
+                [0, np.cos(angles[0]), -np.sin(angles[0])],
+                [0, np.sin(angles[0]), np.cos(angles[0])],
+            ]
+        )
+        Ry = np.array(
+            [
+                [np.cos(angles[1]), 0, np.sin(angles[1])],
+                [0, 1, 0],
+                [-np.sin(angles[1]), 0, np.cos(angles[1])],
+            ]
+        )
+        Rz = np.array(
+            [
+                [np.cos(angles[2]), -np.sin(angles[2]), 0],
+                [np.sin(angles[2]), np.cos(angles[2]), 0],
+                [0, 0, 1],
+            ]
+        )
+        R = np.dot(Rz, np.dot(Ry, Rx))
+        shape_pc = batch_data[k, ...]
+        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
+    return rotated_data
+
+
+def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
+    """Randomly jitter points. jittering is per point.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, jittered batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    assert clip > 0
+    jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1 * clip, clip)
+    jittered_data += batch_data
+    return jittered_data
+
+
+def shift_point_cloud(batch_data, shift_range=0.1):
+    """Randomly shift point cloud. Shift is per point cloud.
+    Input:
+      BxNx3 array, original batch of point clouds
+    Return:
+      BxNx3 array, shifted batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    shifts = np.random.uniform(-shift_range, shift_range, (B, 3))
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] += shifts[batch_index, :]
+    return batch_data
+
+
+def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25):
+    """Randomly scale the point cloud. Scale is per point cloud.
+    Input:
+        BxNx3 array, original batch of point clouds
+    Return:
+        BxNx3 array, scaled batch of point clouds
+    """
+    B, N, C = batch_data.shape
+    scales = np.random.uniform(scale_low, scale_high, B)
+    for batch_index in range(B):
+        batch_data[batch_index, :, :] *= scales[batch_index]
+    return batch_data
+
+
+def random_point_dropout(batch_pc, max_dropout_ratio=0.875):
+    """batch_pc: BxNx3"""
+    for b in range(batch_pc.shape[0]):
+        dropout_ratio = np.random.random() * max_dropout_ratio  # 0~0.875
+        drop_idx = np.where(
+            np.random.random((batch_pc.shape[1])) <= dropout_ratio
+        )[0]
+        if len(drop_idx) > 0:
+            dropout_ratio = (
+                np.random.random() * max_dropout_ratio
+            )  # 0~0.875 # not need
+            batch_pc[b, drop_idx, :] = batch_pc[
+                b, 0, :
+            ]  # set to the first point
+    return batch_pc

examples/pytorch/pointcloud/pointnet/train_cls.py

@@ -0,0 +1,179 @@
+import argparse
+import os
+import urllib
+from functools import partial
+
+import dgl
+
+import provider
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+from dgl.data.utils import download, get_download_dir
+from ModelNetDataLoader import ModelNetDataLoader
+from pointnet2 import PointNet2MSGCls, PointNet2SSGCls
+from pointnet_cls import PointNetCls
+from torch.utils.data import DataLoader
+
+torch.backends.cudnn.enabled = False
+
+
+# from dataset import ModelNet
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", type=str, default="pointnet")
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=200)
+parser.add_argument("--num-workers", type=int, default=8)
+parser.add_argument("--batch-size", type=int, default=32)
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+data_filename = "modelnet40_normal_resampled.zip"
+download_path = os.path.join(get_download_dir(), data_filename)
+local_path = args.dataset_path or os.path.join(
+    get_download_dir(), "modelnet40_normal_resampled"
+)
+
+if not os.path.exists(local_path):
+    download(
+        "https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip",
+        download_path,
+        verify_ssl=False,
+    )
+    from zipfile import ZipFile
+
+    with ZipFile(download_path) as z:
+        z.extractall(path=get_download_dir())
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    loss_f = nn.CrossEntropyLoss()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label in tq:
+            data = data.data.numpy()
+            data = provider.random_point_dropout(data)
+            data[:, :, 0:3] = provider.random_scale_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.jitter_point_cloud(data[:, :, 0:3])
+            data[:, :, 0:3] = provider.shift_point_cloud(data[:, :, 0:3])
+            data = torch.tensor(data)
+            label = label[:, 0]
+
+            num_examples = label.shape[0]
+            data, label = data.to(dev), label.to(dev).squeeze().long()
+            opt.zero_grad()
+            logits = net(data)
+            loss = loss_f(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            num_batches += 1
+            count += num_examples
+            loss = loss.item()
+            correct = (preds == label).sum().item()
+            total_loss += loss
+            total_correct += correct
+
+            tq.set_postfix(
+                {
+                    "AvgLoss": "%.5f" % (total_loss / num_batches),
+                    "AvgAcc": "%.5f" % (total_correct / count),
+                }
+            )
+    scheduler.step()
+
+
+def evaluate(net, test_loader, dev):
+    net.eval()
+
+    total_correct = 0
+    count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label in tq:
+                label = label[:, 0]
+                num_examples = label.shape[0]
+                data, label = data.to(dev), label.to(dev).squeeze().long()
+                logits = net(data)
+                _, preds = logits.max(1)
+
+                correct = (preds == label).sum().item()
+                total_correct += correct
+                count += num_examples
+
+                tq.set_postfix({"AvgAcc": "%.5f" % (total_correct / count)})
+
+    return total_correct / count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if args.model == "pointnet":
+    net = PointNetCls(40, input_dims=6)
+elif args.model == "pointnet2_ssg":
+    net = PointNet2SSGCls(40, batch_size, input_dims=6)
+elif args.model == "pointnet2_msg":
+    net = PointNet2MSGCls(40, batch_size, input_dims=6)
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+opt = optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4)
+
+scheduler = optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.7)
+
+train_dataset = ModelNetDataLoader(local_path, 1024, split="train")
+test_dataset = ModelNetDataLoader(local_path, 1024, split="test")
+train_loader = torch.utils.data.DataLoader(
+    train_dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    num_workers=num_workers,
+    drop_last=True,
+)
+test_loader = torch.utils.data.DataLoader(
+    test_dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=num_workers,
+    drop_last=True,
+)
+
+best_test_acc = 0
+
+for epoch in range(args.num_epochs):
+    train(net, opt, scheduler, train_loader, dev)
+    if (epoch + 1) % 1 == 0:
+        print("Epoch #%d Testing" % epoch)
+        test_acc = evaluate(net, test_loader, dev)
+        if test_acc > best_test_acc:
+            best_test_acc = test_acc
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print("Current test acc: %.5f (best: %.5f)" % (test_acc, best_test_acc))

examples/pytorch/pointcloud/pointnet/train_partseg.py

@@ -0,0 +1,315 @@
+import argparse
+import os
+import time
+import urllib
+from functools import partial
+
+import dgl
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+from dgl.data.utils import download, get_download_dir
+from pointnet2_partseg import PointNet2MSGPartSeg, PointNet2SSGPartSeg
+from pointnet_partseg import PartSegLoss, PointNetPartSeg
+from ShapeNet import ShapeNet
+from torch.utils.data import DataLoader
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", type=str, default="pointnet")
+parser.add_argument("--dataset-path", type=str, default="")
+parser.add_argument("--load-model-path", type=str, default="")
+parser.add_argument("--save-model-path", type=str, default="")
+parser.add_argument("--num-epochs", type=int, default=250)
+parser.add_argument("--num-workers", type=int, default=4)
+parser.add_argument("--batch-size", type=int, default=16)
+parser.add_argument("--tensorboard", action="store_true")
+args = parser.parse_args()
+
+num_workers = args.num_workers
+batch_size = args.batch_size
+
+
+def collate(samples):
+    graphs, cat = map(list, zip(*samples))
+    return dgl.batch(graphs), cat
+
+
+CustomDataLoader = partial(
+    DataLoader,
+    num_workers=num_workers,
+    batch_size=batch_size,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+def train(net, opt, scheduler, train_loader, dev):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.train()
+
+    total_loss = 0
+    num_batches = 0
+    total_correct = 0
+    count = 0
+    start = time.time()
+    with tqdm.tqdm(train_loader, ascii=True) as tq:
+        for data, label, cat in tq:
+            num_examples = data.shape[0]
+            data = data.to(dev, dtype=torch.float)
+            label = label.to(dev, dtype=torch.long).view(-1)
+            opt.zero_grad()
+            cat_ind = [category_list.index(c) for c in cat]
+            # An one-hot encoding for the object category
+            cat_tensor = (
+                torch.tensor(eye_mat[cat_ind])
+                .to(dev, dtype=torch.float)
+                .repeat(1, 2048)
+            )
+            cat_tensor = cat_tensor.view(num_examples, -1, 16).permute(0, 2, 1)
+            logits = net(data, cat_tensor)
+            loss = L(logits, label)
+            loss.backward()
+            opt.step()
+
+            _, preds = logits.max(1)
+
+            count += num_examples * 2048
+            loss = loss.item()
+            total_loss += loss
+            num_batches += 1
+            correct = (preds.view(-1) == label).sum().item()
+            total_correct += correct
+
+            AvgLoss = total_loss / num_batches
+            AvgAcc = total_correct / count
+
+            tq.set_postfix(
+                {"AvgLoss": "%.5f" % AvgLoss, "AvgAcc": "%.5f" % AvgAcc}
+            )
+    scheduler.step()
+    end = time.time()
+    return data, preds, AvgLoss, AvgAcc, end - start
+
+
+def mIoU(preds, label, cat, cat_miou, seg_classes):
+    for i in range(preds.shape[0]):
+        shape_iou = 0
+        n = len(seg_classes[cat[i]])
+        for cls in seg_classes[cat[i]]:
+            pred_set = set(np.where(preds[i, :] == cls)[0])
+            label_set = set(np.where(label[i, :] == cls)[0])
+            union = len(pred_set.union(label_set))
+            inter = len(pred_set.intersection(label_set))
+            if union == 0:
+                shape_iou += 1
+            else:
+                shape_iou += inter / union
+        shape_iou /= n
+        cat_miou[cat[i]][0] += shape_iou
+        cat_miou[cat[i]][1] += 1
+
+    return cat_miou
+
+
+def evaluate(net, test_loader, dev, per_cat_verbose=False):
+    category_list = sorted(list(shapenet.seg_classes.keys()))
+    eye_mat = np.eye(16)
+    net.eval()
+
+    cat_miou = {}
+    for k in shapenet.seg_classes.keys():
+        cat_miou[k] = [0, 0]
+    miou = 0
+    count = 0
+    per_cat_miou = 0
+    per_cat_count = 0
+
+    with torch.no_grad():
+        with tqdm.tqdm(test_loader, ascii=True) as tq:
+            for data, label, cat in tq:
+                num_examples = data.shape[0]
+                data = data.to(dev, dtype=torch.float)
+                label = label.to(dev, dtype=torch.long)
+                cat_ind = [category_list.index(c) for c in cat]
+                cat_tensor = (
+                    torch.tensor(eye_mat[cat_ind])
+                    .to(dev, dtype=torch.float)
+                    .repeat(1, 2048)
+                )
+                cat_tensor = cat_tensor.view(num_examples, -1, 16).permute(
+                    0, 2, 1
+                )
+                logits = net(data, cat_tensor)
+                _, preds = logits.max(1)
+
+                cat_miou = mIoU(
+                    preds.cpu().numpy(),
+                    label.view(num_examples, -1).cpu().numpy(),
+                    cat,
+                    cat_miou,
+                    shapenet.seg_classes,
+                )
+                for _, v in cat_miou.items():
+                    if v[1] > 0:
+                        miou += v[0]
+                        count += v[1]
+                        per_cat_miou += v[0] / v[1]
+                        per_cat_count += 1
+                tq.set_postfix(
+                    {
+                        "mIoU": "%.5f" % (miou / count),
+                        "per Category mIoU": "%.5f" % (miou / count),
+                    }
+                )
+    if per_cat_verbose:
+        print("Per-Category mIoU:")
+        for k, v in cat_miou.items():
+            if v[1] > 0:
+                print("%s mIoU=%.5f" % (k, v[0] / v[1]))
+            else:
+                print("%s mIoU=%.5f" % (k, 1))
+    return miou / count, per_cat_miou / per_cat_count
+
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# dev = "cpu"
+if args.model == "pointnet":
+    net = PointNetPartSeg(50, 3, 2048)
+elif args.model == "pointnet2_ssg":
+    net = PointNet2SSGPartSeg(50, batch_size, input_dims=6)
+elif args.model == "pointnet2_msg":
+    net = PointNet2MSGPartSeg(50, batch_size, input_dims=6)
+
+net = net.to(dev)
+if args.load_model_path:
+    net.load_state_dict(
+        torch.load(args.load_model_path, weights_only=False, map_location=dev)
+    )
+
+opt = optim.Adam(net.parameters(), lr=0.001, weight_decay=1e-4)
+scheduler = optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.5)
+L = PartSegLoss()
+
+shapenet = ShapeNet(2048, normal_channel=False)
+
+train_loader = CustomDataLoader(shapenet.trainval())
+test_loader = CustomDataLoader(shapenet.test())
+
+# Tensorboard
+if args.tensorboard:
+    import torchvision
+    from torch.utils.tensorboard import SummaryWriter
+    from torchvision import datasets, transforms
+
+    writer = SummaryWriter()
+# Select 50 distinct colors for different parts
+color_map = torch.tensor(
+    [
+        [47, 79, 79],
+        [139, 69, 19],
+        [112, 128, 144],
+        [85, 107, 47],
+        [139, 0, 0],
+        [128, 128, 0],
+        [72, 61, 139],
+        [0, 128, 0],
+        [188, 143, 143],
+        [60, 179, 113],
+        [205, 133, 63],
+        [0, 139, 139],
+        [70, 130, 180],
+        [205, 92, 92],
+        [154, 205, 50],
+        [0, 0, 139],
+        [50, 205, 50],
+        [250, 250, 250],
+        [218, 165, 32],
+        [139, 0, 139],
+        [10, 10, 10],
+        [176, 48, 96],
+        [72, 209, 204],
+        [153, 50, 204],
+        [255, 69, 0],
+        [255, 145, 0],
+        [0, 0, 205],
+        [255, 255, 0],
+        [0, 255, 0],
+        [233, 150, 122],
+        [220, 20, 60],
+        [0, 191, 255],
+        [160, 32, 240],
+        [192, 192, 192],
+        [173, 255, 47],
+        [218, 112, 214],
+        [216, 191, 216],
+        [255, 127, 80],
+        [255, 0, 255],
+        [100, 149, 237],
+        [128, 128, 128],
+        [221, 160, 221],
+        [144, 238, 144],
+        [123, 104, 238],
+        [255, 160, 122],
+        [175, 238, 238],
+        [238, 130, 238],
+        [127, 255, 212],
+        [255, 218, 185],
+        [255, 105, 180],
+    ]
+)
+
+
+# paint each point according to its pred
+def paint(batched_points):
+    B, N = batched_points.shape
+    colored = color_map[batched_points].squeeze(2)
+    return colored
+
+
+best_test_miou = 0
+best_test_per_cat_miou = 0
+
+for epoch in range(args.num_epochs):
+    data, preds, AvgLoss, AvgAcc, training_time = train(
+        net, opt, scheduler, train_loader, dev
+    )
+    if (epoch + 1) % 5 == 0:
+        print("Epoch #%d Testing" % epoch)
+        test_miou, test_per_cat_miou = evaluate(
+            net, test_loader, dev, (epoch + 1) % 5 == 0
+        )
+        if test_miou > best_test_miou:
+            best_test_miou = test_miou
+            best_test_per_cat_miou = test_per_cat_miou
+            if args.save_model_path:
+                torch.save(net.state_dict(), args.save_model_path)
+        print(
+            "Current test mIoU: %.5f (best: %.5f), per-Category mIoU: %.5f (best: %.5f)"
+            % (
+                test_miou,
+                best_test_miou,
+                test_per_cat_miou,
+                best_test_per_cat_miou,
+            )
+        )
+    # Tensorboard
+    if args.tensorboard:
+        colored = paint(preds)
+        writer.add_mesh(
+            "data", vertices=data, colors=colored, global_step=epoch
+        )
+        writer.add_scalar(
+            "training time for one epoch", training_time, global_step=epoch
+        )
+        writer.add_scalar("AvgLoss", AvgLoss, global_step=epoch)
+        writer.add_scalar("AvgAcc", AvgAcc, global_step=epoch)
+        if (epoch + 1) % 5 == 0:
+            writer.add_scalar("test mIoU", test_miou, global_step=epoch)
+            writer.add_scalar(
+                "best test mIoU", best_test_miou, global_step=epoch
+            )