torch-pointcloud: pretrained point cloud models (PointNet++, KPConv, PTv3...) built on PyG

[arthurdjn]

Hi PyG team and community! I'd like to share torch-pointcloud, a library of pretrained point cloud models built on PyG: one create_model call loads an architecture and its weights, in the spirit of timm.

import torch
import torch_pointcloud as tp

model = tp.create_model(
  "pointnext-sm.scanobjectnn-hardest.openpoints", 
  task="classification",
  pretrained=True,
).eval()

pos = torch.randn(2048, 3)
x = torch.cat([pos, pos[:, 1:2] - pos[:, 1].min()], dim=1)  # xyz + height
batch = torch.zeros(2048, dtype=torch.long)                 # PyG packed-batch convention
logits = model(x, pos, batch)                               # (1, 15)

How it uses PyG

• Every model takes PyG's packed format: a flat (N, ...) tensor plus a batch vector, never padded tensors.
• Sampling and neighbor search (FPS, kNN, radius) go through pyg-lib.
• Most convolutions are built on MessagePassing.

What's in it: 30+ architectures and associated pretrained checkpoints for classification, segmentation, detection and self-supervised pretraining, plus datasets, transforms and tiled inference for large scenes. The weights come from their original authors (Pointcept, OpenPoints, MIT Han Lab, FAIR, OpenMMLab and others), and each checkpoint lists its source and license.

It's still alpha, so feedback is very welcome, especially on the PyG side: anything I'm using in a non-idiomatic way, or PyG features I should rely on more.

• Repo: https://github.com/arthurdjn/pytorch-pointcloud
• Docs: https://pytorch-pointcloud.org