Model card for litept-small.scannet20.prs-eth

A LitePT point cloud semantic segmentation model (convolutions in the early stages, attention with 3D rotary embedding in the late ones). Trained on ScanNet (20 classes).

Model Details

Install

pip install torch-pointcloud

This checkpoint also needs spconv and flash-attn, which need a build matching your torch and CUDA: see the installation guide.

Usage

import torch
import torch_pointcloud as tp
from torch_pointcloud.utils.data import collate

model, info = tp.create_model(
    "litept-small.scannet20.prs-eth",
    pretrained=True,
    return_info=True,
)
model = model.cuda().eval()  # GPU-only kernels

# synthetic sample with the keys a dataset provides
num_points = 8192
sample = {
    "pos": torch.randn(num_points, 3),
    "color": torch.rand(num_points, 3) * 255,
    "normal": torch.randn(num_points, 3),
    "segment": torch.zeros(num_points, dtype=torch.long),
    "instance": torch.zeros(num_points, dtype=torch.long),
}
data = info["transform"](sample)
data = collate([data])
data = {key: value.cuda() for key, value in data.items()}

with torch.no_grad():
    logits = model(data.get("x"), data["pos_grid"], data["batch"])

Feature extraction

with torch.no_grad():
    features = model.forward_features(data.get("x"), data["pos_grid"], data["batch"])

model.reset_classifier(num_classes=0)
with torch.no_grad():
    features = model(data.get("x"), data["pos_grid"], data["batch"])  # (N, 72)

Citation

@inproceedings{yue2026litept,
  title   = {LitePT: Lighter Yet Stronger Point Transformer},
  author  = {Yuanwen Yue and Damien Robert and Jianyuan Wang and Sunghwan Hong and Jan Dirk Wegner and Christian Rupprecht and Konrad Schindler},
  booktitle = {CVPR},
  year    = {2026}
}

@inproceedings{dai2017scannet,
  title   = {ScanNet: Richly-annotated 3D Reconstructions of Indoor Scenes},
  author  = {Angela Dai and Angel X. Chang and Manolis Savva and Maciej Halber and Thomas Funkhouser and Matthias Nießner},
  booktitle = {CVPR},
  year    = {2017}
}

@software{dujardin2026pytorchpointcloud,
  author  = {Arthur Dujardin},
  title   = {PyTorch PointCloud},
  year    = {2026},
  doi     = {10.5281/zenodo.22159632},
  url     = {https://github.com/arthurdjn/pytorch-pointcloud},
}
Downloads last month
32
Safetensors
Model size
12.7M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for torch-pointcloud/litept-small.scannet20.prs-eth

Evaluation results