torch-pointcloud：遵循 timm 风格的 3D 点云预训练模型 PyTorch 库

https://preview.redd.it/m5e9anluo8vh1.png?width=1080&format=png&auto=webp&s=4c80687e6fe40036843a97ed8975ce46eaab0662

I am excited to share a project I have been working on during the past few years during my free time. I started this project after being unable to use torch-points3d at the time (too many installation errors).

Note that this project is still in active development and in alpha. I am looking for early users and feedback!

What it is: torch-pointcloud is a plain-PyTorch library that puts multiple point cloud architectures and their pretrained checkpoints behind one create_model() factory, in the spirit of timm.

This is a library that focuses on the reusability and standardization of core components and layers, i.e. KNN, FPS and other PointNetSetAbstraction are not re-implemented over and over. In this library, all models share the same interface and reuse the same core components.

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) logits = model(x, pos, batch) # (1, 15) 

Install (see the docs for more details):

# Optional: specify your CUDA version for torch pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cu128 pip install torch-pointcloud # Install required CUDA/CPU kernels associated with your CUDA version pip install pyg-lib -f https://data.pyg.org/whl/torch-2.10.0+cu128.html 

I tried to make the installation process smoother by not maintaining CUDA kernels myself (not enough time currently, so I tried to reuse already pre-built wheels as much as I could).

Scope:

• Classification, semantic and part segmentation, detection, and self-supervised pretraining: PointNet++, DGCNN, KPConv, RandLA-Net, PointNeXt, PointMLP, Point Transformer v1/v2/v3, Sonata, PointMamba, VoteNet, 3DETR, PointPillars, SECOND, PointRCNN, SPVCNN, OctFormer, Point-MAE/BERT/GPT and more.
• Weights are hosted on Hugging Face.
• Datasets (ModelNet40, ScanObjectNN, S3DIS, ScanNet, SemanticKITTI, nuScenes, KITTI, SUN RGB-D...), transforms, tiled inferers for large scale segmentation.
• Packed-batch tensors (the PyTorch Geometric convention).

Supported tasks and use cases

But... why?

Most point cloud libraries live in research codebases built around configs and custom ops and kernels to build locally. It is also sometimes difficult to reproduce past released experiments from research codebases (typically the code associated with the checkpoint changed with no notice). In addition some models are evaluated using different strategies (TTA, partitioning, etc.) which makes comparison unfair as they are not evaluated using a common strategy.

In contrast, I wanted something I could pip install easily, reuse weights trained from different repos and benchmark easily.

All models were re-benchmarked and compared against original reported metrics.

Measured performances vs original ones

Credits: the architectures and pretrained weights come from their original authors, notably Pointcept, OpenPoints, MIT Han Lab, FAIR, OpenMMLab and more. Every checkpoint name ends with its source author (e.g. .pointcept). The THIRD_PARTY_NOTICES.md lists every source repo, license and dataset term.

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