FanMingyu/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.