cstnet2-stage1-pointnet2-baseline

CSTNet2 Stage 1 PointNet++ Direct Baseline

This repository contains the PointNet++ direct-prediction baseline used for comparison with CSTNet2 Stage 1. Given XYZ points, one shared PointNet++ backbone and four independent heads directly predict primitive type, direction, dimension, and location. This baseline does not use instance embedding, clustering, or geometric primitive fitting.

Source code: xcheng-tsinghua/cstnet2

Training dataset: ZXCCHENGXI/cstnet2_stage1_mini

Files

• best_pmt_miou.pth: recommended checkpoint for primitive prediction comparison
• best_loss.pth: checkpoint selected by total direct-prediction training loss
• last.pth: final epoch checkpoint
• history.json: all 80 training epoch summaries
• config.json: architecture and training configuration
• checksums.sha256: SHA256 verification values

Training configuration

• Backbone: PointNet++
• 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

Training-set metrics

• Best primitive mIoU: 0.547805
• Best total loss: 867.963349

These values are training-set metrics. No held-out validation split was used for checkpoint selection.

Loading

import torch

from networks.stage1_direct_baselines import build_stage1_direct_baseline

checkpoint = torch.load(
    "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()

Limitations

• Metrics above are training metrics and should not be interpreted as held-out generalization results.
• This is a direct regression baseline, not a replacement for the full clustering-and-fitting CSTNet2 Stage 1 pipeline.
• Checkpoint files use Python/PyTorch serialization. Load only files from trusted sources.