CSTNet2 Stage 1 Backbone Baselines
This repository contains eight independently trained XYZ-only Stage 1 direct baselines. Each model uses one point-cloud backbone and four per-point heads to predict primitive type, direction, dimension, and location. The baselines do not load CSTNet2 Stage 1 weights and do not use instance embedding, clustering, or geometric primitive fitting.
Source code: xcheng-tsinghua/cstnet2
Training dataset: ZXCCHENGXI/cstnet2_stage1_mini
Training configuration
- Points per sample: 2,048
- Epochs: 80
- Hardware: 4 GPUs with DDP
- Precision: BF16 AMP
- Per-GPU batch size: 32 (global batch size 128)
- Seed: 2026
- Learning rate: 1e-4
- Weight decay: 1e-4
All eight models were randomly initialized and trained independently with the same dataset and optimization settings.
Training-set leaderboard
| Backbone | Best primitive mIoU | Best total loss |
|---|---|---|
| PointMLP | 0.654248 | 867.811967 |
| DGCNN | 0.570031 | 867.934994 |
| PointNet++ | 0.547805 | 867.963349 |
| PointNeXt | 0.545523 | 867.983210 |
| PointNet | 0.531966 | 867.991856 |
| Attn3DGCN | 0.448774 | 867.927836 |
| PointMamba | 0.391190 | 868.047581 |
| PointTransformer | 0.385701 | 868.044473 |
These values are training-set metrics. No held-out validation split was used for checkpoint selection, so they must not be interpreted as generalization results.
Repository layout
<backbone>/seed_2026/
best_pmt_miou.pth
best_loss.pth
last.pth
history.json
config.json
The available backbone names are pointnet2, pointnet, attn3dgcn,
dgcnn, pointtransformer, pointmamba, pointnext, and pointmlp.
Loading
import torch
from networks.stage1_direct_baselines import build_stage1_direct_baseline
checkpoint = torch.load(
"pointmlp/seed_2026/best_pmt_miou.pth",
map_location="cpu",
weights_only=False,
)
model = build_stage1_direct_baseline(checkpoint["model_config"])
model.load_state_dict(checkpoint["model"], strict=True)
model.eval()
Related repository
The original standalone PointNet++ upload remains available at FanMingyu/cstnet2-stage1-pointnet2-baseline.
Limitations
- Reported scores are training metrics only.
- The models were trained on mechanical CAD-derived point clouds and may not generalize to noisy real-world scans.
- Checkpoint files use Python/PyTorch serialization. Load only files from trusted sources.