PyTorch implementation of SQN (Semantic Query Network) for weakly-supervised 3D point cloud semantic segmentation.
Train 3D semantic segmentation with just 0.1% labeled points. SQN replaces RandLA-Net's decoder with a Semantic Query Module that propagates sparse annotations to dense predictions.
- Modern PyTorch (CUDA 11.x/12.x compatible)
- First PyTorch reproduction (original: TensorFlow 1.x)
- Class-stratified pseudo-labeling
- Vote-based evaluation matching official protocol
- Pre-trained checkpoints
| Method | Labels | mIoU (%) | Notes |
|---|---|---|---|
| RandLA-Net (Baseline) | 100% | 62.59 | Full supervision |
| SQN (Ours) | 0.1% | 57.52 | Weakly-supervised |
| SQN + Pseudo-Label | 0.1% | 60.55 | +3.03% improvement |
| SQN (Paper, TensorFlow) | 0.1% | 61.4 | Gap: -0.85% |
| Method | mIoU (%) | Year |
|---|---|---|
| SQN + Pseudo-Label (Ours) | 60.55 | - |
| SQN (Paper) | 61.4 | ECCV 2022 |
| CPCM (MinkNet) | 66.3 | CVPR 2024 |
| PointCT (PointTrans) | 71.2 | WACV 2024 |
Requirements: Python 3.7, PyTorch 1.12.1+, CUDA 11.x
git clone https://github.com/yourusername/SQN-pytorch.git
cd SQN-pytorch
pip install -r requirements.txt
bash compile_op.sh # Compile C++ extensions (KNN, grid subsampling)-
Download S3DIS from Stanford
-
Create symlink:
ln -s /path/to/Stanford3dDataset_v1.2_Aligned_Version data/S3DIS
-
Preprocess:
python utils/data_prepare_s3dis.py
See data/README.md for details.
# Basic (0.1% labels)
python train_scripts/main_SQN_S3DIS.py --test_area 5 --gpu 0 --labeled_point 0.1% --max_epoch 100
# With wandb logging
python train_scripts/main_SQN_S3DIS.py --test_area 5 --gpu 0 --labeled_point 0.1% --use_wandb
# Different label ratios
python train_scripts/main_SQN_S3DIS.py --test_area 5 --gpu 0 --labeled_point 1%
python train_scripts/main_SQN_S3DIS.py --test_area 5 --gpu 0 --labeled_point 1 # 1 point/class# Standard
python test_scripts/test_SQN_S3DIS.py --checkpoint_path path/to/checkpoint.tar --test_area 5 --labeled_point 0.1%
# Generate pseudo-labels
python test_scripts/test_SQN_S3DIS.py --checkpoint_path path/to/checkpoint.tar --test_area 5 --labeled_point 0.1% --gen_pseudo --pseudo_dir pseudo_labels# Step 1: Generate pseudo-labels
python test_scripts/test_SQN_S3DIS.py --checkpoint_path path/to/best_checkpoint.tar --test_area 5 --labeled_point 0.1% --gen_pseudo --pseudo_dir pseudo_labels --pl_strategy stratified
# Step 2: Retrain
python train_scripts/main_SQN_S3DIS.py --test_area 5 --gpu 0 --labeled_point 0.1% --retrain --pseudo_dir pseudo_labels# Train
python train_scripts/main_S3DIS.py --test_area 5 --gpu 0 --max_epoch 100
# Test
python test_scripts/test_S3DIS.py --checkpoint_path path/to/checkpoint.tar --test_area 5Download all checkpoints from Google Drive.
| Model | Labels | mIoU |
|---|---|---|
| SQN + Pseudo-Label | 0.1% | 60.55% |
| SQN | 0.1% | 57.52% |
| RandLA-Net | 100% | 62.59% |
python test_scripts/test_SQN_S3DIS.py --checkpoint_path sqn_s3dis_0.1pct_pseudo.tar --test_area 5 --labeled_point 0.1%Input (XYZRGB) → fc0 (6→8)
↓
Encoder (5 Dilated Residual Blocks + Random Sampling)
Layer 1: 8→32 dims, 40960→10240 pts
Layer 2: 32→128 dims, 10240→2560 pts
Layer 3: 128→256 dims, 2560→640 pts
Layer 4: 256→512 dims, 640→160 pts
Layer 5: 512→1024 dims, 160→40 pts
↓
Semantic Query Module (replaces decoder)
KNN interpolation (K=3) from each encoder scale
Concatenate multi-scale features: 1952 dims
MLP: 1952→256→128→64→64→32→13
↓
Output: Per-point semantic labels
Key Differences from RandLA-Net:
| RandLA-Net | SQN |
|---|---|
| Decoder with upsampling | Semantic Query Module |
| All points as queries | Sparse annotated points (M≈40) during training |
| Standard batch size | Augmentation doubling for effective batch |
SQN-pytorch/
├── sqn/ # SQN core
│ ├── SQN_modules.py # Semantic Query Module
│ └── sqn_s3dis_dataset.py
├── randlanet/ # Base RandLA-Net
│ ├── RandLANet.py
│ ├── helper_tool.py
│ └── helper_ply.py
├── datasets/ # Dataset loaders
├── train_scripts/ # Training scripts
├── test_scripts/ # Testing scripts
├── utils/ # C++ extensions & preprocessing
├── requirements.txt
└── LICENSE
@inproceedings{hu2022sqn,
title={SQN: Weakly-Supervised Semantic Segmentation of Large-Scale 3D Point Clouds},
author={Hu, Qingyong and Fang, Bo and Wang, Anquan and Liu, Shengjia and Guo, Yulan and Wang, Qingjie and Liu, Li and Chen, Chengzhong and Wang, Nenglun and Cheng, Ming-Ming},
booktitle={ECCV},
pages={640--657},
year={2022}
}@article{hu2020randla,
title={RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds},
author={Hu, Qingyong and Yang, Bo and Xie, Linhai and Rosa, Stefano and Guo, Yulan and Wang, Zhihua and Trigoni, Niki and Markham, Andrew},
journal={CVPR},
year={2020}
}- QingyongHu/SQN — Official TensorFlow implementation
- QingyongHu/RandLA-Net — RandLA-Net base
- qiqihaer/RandLA-Net-pytorch — PyTorch reference
MIT License — see LICENSE.