ScanNetV2 Benchmark: AQ3D achieves mAP: 65.6 and AP@50: 83.4 (screenshots from August 26, 2026)
ScanNet200 Benchmark: AQ3D achieves mAP: 38.5 and AP@50: 49.1 (screenshots from August 26, 2026)
ScanNet++V2 Benchmark: AQ3D achieves mAP: 38.0 and AP@50: 55.7 (screenshots from August 26, 2026)
| Validation | Test | Weights | ||||
|---|---|---|---|---|---|---|
| Dataset | Backbone | mAP | AP@50 | mAP | AP@50 | |
| ScanNetV2 | Sparse UNet (44M) | 65.6 | 82.1 | 65.6 | 83.4 | aq3d_scannet_spunet.pth |
| ScanNet200 | Sparse UNet (44M) | 39.6 | 49.7 | - | - | aq3d_scannet200_spunet.pth |
| Volt-B (88M) | 44.1 | 55.1 | 38.5 | 49.1 | aq3d_scannet200_volt.pth | |
| ScanNet++V2 | Sparse UNet (44M) | 36.9 | 53.2 | - | - | aq3d_scannetpp_spunet.pth |
| Volt-B (88M) | 38.3 | 57.1 | 38.0 | 55.7 | aq3d_scannetpp_volt.pth | |
We provide a ready-to-use Dockerfile with all required dependencies (CUDA 12.8, PyTorch 2.7.1, spconv, FlashAttention, Pointcept, segmentator, etc.) together with a matching devcontainer configuration for VS Code. If you have the Dev Containers extension and the NVIDIA Container Toolkit installed, simply open the repository in VS Code and select "Reopen in Container" — this builds the image and drops you into a fully configured, GPU-enabled environment. Alternatively, build and run the image manually:
docker build -t aq3d -f docker/Dockerfile .
docker run --gpus all --ipc=host -v "$(pwd)":/workspace -w /workspace -it aq3d bashIf you prefer a native (non-Docker) setup, install the Python dependencies listed in requirements.txt (or requirements_frozen.txt for the exact, pinned versions we used) yourself, alongside the system/CUDA libraries listed in the Dockerfile.
This repository depends on the ScanNet toolkit and Pointcept, which are included as git submodules. After cloning the repository, fetch them with:
git submodule update --initAQ3D is trained and evaluated on ScanNetV2/ScanNet200 and ScanNet++V2. Both datasets require registering with the respective dataset authors to obtain download access, after which the preprocessing scripts below convert the raw scans into the point cloud format (coordinates, colors, normals, semantic/instance labels) expected by the dataloaders, and compute an over-segmentation with Segmentator that is used for superpoint pooling/grouping.
Copy the raw dataset into data/scannet/raw and run the commands below. The first pair of commands preprocesses the standard 20-class ScanNetV2 benchmark into data/scannet/processed; the second pair preprocesses the same scans with the 200-class ScanNet200 label set into data/scannet/processed200. In both cases, segmentator is run afterwards on the preprocessed data to compute the superpoints:
python -m datasets.scannet.scannet_3d_preprocessing preprocess \
--raw_dir="./data/scannet/raw" --processed_dir="./data/scannet/processed"
python -m datasets.segmentator preprocess \
--raw_dir="./data/scannet/processed" --processed_dir="./data/scannet/processed"python -m datasets.scannet.scannet_3d_preprocessing preprocess --scannet200=True \
--raw_dir="./data/scannet/raw" --processed_dir="./data/scannet/processed200"
python -m datasets.segmentator preprocess \
--raw_dir="./data/scannet/processed200" --processed_dir="./data/scannet/processed200"Copy the raw dataset into data/scannetpp/raw and run the commands below. As with ScanNet, scans are first preprocessed into data/scannetpp/processed and segmentator is run on them to compute superpoints. Since ScanNet++ scenes are considerably larger than ScanNet scenes, they are additionally split into overlapping spatial chunks with datasets.chunk, written to data/scannetpp/processed_train_chunked, which is the data split actually used for training:
python -m datasets.scannetpp.scannetpp_3d_preprocessing preprocess \
--raw_dir="./data/scannetpp/raw" --processed_dir="./data/scannetpp/processed"
python -m datasets.segmentator preprocess \
--raw_dir="./data/scannetpp/processed" --processed_dir="./data/scannetpp/processed" \
--kThresh=0.2 --segMinVerts=100
python -m datasets.chunk preprocess \
--raw_dir="./data/scannetpp/processed" --processed_dir="./data/scannetpp/processed_train_chunked" \
--chunk_range="(6, 6)" --chunk_stride="(3, 3)" --chunk_minimum_size=10000
cp ./data/scannetpp/processed/label_database.yaml ./data/scannetpp/processed_train_chunked/label_database.yamlWithin this work, by default only the training split is chunked, while the validation/test splits are copied as-is so that evaluation is still performed on full, unchunked scenes. When you also want to chunk the validation split (set --modes and --copy accordingly):
python -m datasets.chunk preprocess \
--raw_dir="./data/scannetpp/processed" --processed_dir="./data/scannetpp/processed_train_val_chunked" \
--chunk_range="(6, 6)" --chunk_stride="(3, 3)" --chunk_minimum_size=10000 \
--modes="('train', 'val', 'test')" --copy="('test')"
cp ./data/scannetpp/processed/label_database.yaml ./data/scannetpp/processed_train_val_chunked/label_database.yamlTraining is driven by Hydra configs under configs and launched via src/train.py. Below we describe how to obtain pretrained backbone weights, the available AQ3D model/data/transform combinations, and how to train the other models (SPFormer, MAFT, Relation3D), and the sparse UNet backbone for semantic segmentation, which we ship with this codebase.
AQ3D uses either a sparse UNet or Volt as its point cloud backbone. Both are initialized from pretrained weights before AQ3D training rather than trained from scratch; download the weights relevant to the backbone/dataset combination you want to train, as described below.
Official Volt backbone weights, released by the Volt authors:
mkdir -p weights/volt
curl -L -o weights/volt/volt-base-scannet200.pth https://huggingface.co/KadirYilmaz/Volt/resolve/main/Volt_experiments/joint_training_base/scannet200/model/model_last.pth
curl -L -o weights/volt/volt-base-scannetpp.pth https://huggingface.co/KadirYilmaz/Volt/resolve/main/Volt_experiments/joint_training_base/scannetpp/model/model_last.pthThe sparse UNet checkpoint pretrained by SSTNet, used to initialize the sparse UNet backbone when training SPFormer, MAFT, and Relation3D, following those methods' official setup:
curl -L -o weights/sstnet_pretrain.pth "https://drive.usercontent.google.com/download?id=1vucwdbm6pHRGlUZAYFdK9JmnPVerjNuD&export=download&confirm=t"The sparse UNet backbone, pretrained via semantic segmentation on ScanNet, used to initialize the sparse UNet backbone for AQ3D training:
curl -L -o weights/sparse_unet.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/sparse_unet.pthThe table below lists all AQ3D dataset/backbone combinations we support, along with the corresponding Hydra model, data/datasets, and data/transforms config names to substitute into the training command:
| Dataset | Backbone | Batch size | Epochs | Model | Data | Transforms |
|---|---|---|---|---|---|---|
| ScanNetV2 | Sparse UNet | 4 | 512 | aqtd_scannet |
scannet_instance_coord_color_normal |
scannet_spformer_like |
| ScanNet200 | Sparse UNet | 4 | 512 | aqtd_scannet200 |
scannet200_instance_coord_color_normal |
scannet_spformer_like |
| ScanNet200 | Volt | 4 | 384 | aqtd_volt_scannet200 |
scannet200_instance_color_normal |
scannet_spformer_like |
| ScanNet++V2 | Sparse UNet | 4 | 512 | aqtd_scannetpp |
scannetpp_instance_coord_color_normal |
scannetpp_spformer_like |
| ScanNet++V2 | Volt | 8 | 384 | aqtd_volt_scannetpp |
scannetpp_instance_color_normal |
scannetpp_volt |
Replace <model>, <data>, <transforms>, <batch_size> and <epochs> with the values from the table above, and <PROJECT-XY>/<RUN-NAME> with your Weights & Biases project and run name to launch training:
python src/train.py model=<model> trainer.max_epochs=<epochs> data.batch_size=<batch_size> \
data/datasets=<data> data/transforms=<transforms> \
logger=wandb logger.wandb.project=<PROJECT-XY> +logger.wandb.name=<RUN-NAME>If your GPU cannot fit the batch size above, use gradient accumulation to reach the same effective batch size at lower memory cost (here, an accumulation of 2 steps at batch size 4 reproduces an effective batch size of 8, e.g. for aqtd_volt_scannetpp):
python src/train.py model=<model> trainer.max_epochs=<epochs> \
data.batch_size=4 +trainer.accumulate_grad_batches=2 \
data/datasets=<data> data/transforms=<transforms> \
logger=wandb logger.wandb.project=<PROJECT-XY> +logger.wandb.name=<RUN-NAME>For larger scenes/batch sizes (e.g. ScanNet++V2), CUDA memory fragmentation can trigger out-of-memory errors despite enough total memory being free. Setting the following environment variable before launching training mitigates this:
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:TrueTo evaluate on the chunked validation split of ScanNet++ during training, point the dataloaders at processed_train_val_chunked instead of the default processed_train_chunked directory:
data.train_dataset.data_dir="data/scannetpp/processed_train_val_chunked"
data.val_dataset.data_dir="data/scannetpp/processed_train_val_chunked"When using Volt as a backbone, add gradient clipping:
+trainer.gradient_clip_val=10.0
+trainer.gradient_clip_algorithm="norm"To obtain the "Sparse UNet for AQ3D" backbone weights referenced above, the sparse UNet is first pretrained on semantic segmentation of superpoints (second following command) on ScanNet. The task includes the superpoint pooling mechanism as used in AQ3D.
python src/train.py model=spunet trainer.max_epochs=256 data.batch_size=8 \
data/datasets=scannet_semantic_coord_color_normal data/transforms=scannet_spformer_like \
logger=csv
python src/train.py model=spunet_superpoint trainer.max_epochs=256 data.batch_size=8 \
data/datasets=scannet_semantic_superpoint_coord_color_normal data/transforms=scannet_spformer_like \
logger=csvReference reimplementation of SPFormer, shipped within this codebase:
python src/train.py model=spformer_like trainer.max_epochs=512 data.batch_size=4 \
data/datasets=scannet_instance_coord_color_normal data/transforms=scannet_spformer_like \
logger=wandb logger.wandb.project=<PROJECT-XY> +logger.wandb.name=<RUN-NAME>Reference reimplementation of MAFT, shipped within this codebase (without attention_rpe_ops overhead):
python src/train.py model=maft_like trainer.max_epochs=512 data.batch_size=4 \
data/datasets=scannet_instance_coord_color_normal data/transforms=scannet_spformer_like \
logger=wandb logger.wandb.project=<PROJECT-XY> +logger.wandb.name=<RUN-NAME>Reference reimplementation of Relation3D, shipped within this codebase (without attention_rpe_ops overhead). This model additionally requires the spatials and spgtindices instance keys to be produced by the dataloader, which is enabled via the flags below:
python src/train.py model=relation3d_like trainer.max_epochs=512 data.batch_size=4 \
data/datasets=scannet_instance_coord_color_normal data/transforms=scannet_spformer_like \
data.train_dataset.instance_keys=['spatials','spgtindices'] \
data.val_dataset.instance_keys=['spatials','spgtindices'] \
logger=wandb logger.wandb.project=<PROJECT-XY> +logger.wandb.name=<RUN-NAME>We release the pretrained AQ3D checkpoints corresponding to the results reported in the AQ3D Results table above, one per dataset/backbone combination. Download the weight file into weights/ (matching the ckpt_path used below), then run the corresponding src/eval.py command to reproduce our reported numbers. data/datasets/data/transforms must match the ones the checkpoint was trained with.
Sparse UNet as backbone:
curl -L -o weights/aq3d_scannet_spunet.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/aq3d_scannet_spunet.pth
python src/eval.py model=aqtd_scannet +weights_only=True ckpt_path='weights/aq3d_scannet_spunet.pth' data/datasets=scannet_instance_coord_color_normal data/transforms=scannet_spformer_likeSparse UNet as backbone:
curl -L -o weights/aq3d_scannet200_spunet.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/aq3d_scannet200_spunet.pth
python src/eval.py model=aqtd_scannet200 +weights_only=True ckpt_path='weights/aq3d_scannet200_spunet.pth' data/datasets=scannet200_instance_coord_color_normal data/transforms=scannet_spformer_likeVolt as backbone:
curl -L -o weights/aq3d_scannet200_volt.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/aq3d_scannet200_volt.pth
python src/eval.py model=aqtd_volt_scannet200 +weights_only=True ckpt_path='weights/aq3d_scannet200_volt.pth' data/datasets=scannet200_instance_color_normal data/transforms=scannet_spformer_likeSparse UNet as backbone:
curl -L -o weights/aq3d_scannetpp_spunet.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/aq3d_scannetpp_spunet.pth
python src/eval.py model=aqtd_scannetpp +weights_only=True ckpt_path='weights/aq3d_scannetpp_spunet.pth' data/datasets=scannetpp_instance_coord_color_normal data/transforms=scannetpp_spformer_like data.val_dataset.data_dir="data/scannetpp/processed_train_val_chunked"Volt as backbone:
curl -L -o weights/aq3d_scannetpp_volt.pth https://huggingface.co/kenomo/aq3d/resolve/main/weights/aq3d_scannetpp_volt.pth
python src/eval.py model=aqtd_volt_scannetpp +weights_only=True ckpt_path='weights/aq3d_scannetpp_volt.pth' data/datasets=scannetpp_instance_color_normal data/transforms=scannetpp_volt data.val_dataset.data_dir="data/scannetpp/processed_train_val_chunked"If you find AQ3D useful to your work/research, please cite:
@misc{Moenck.2026,
title = {AQ3D: Adaptive Query Transformer for 3D Instance Segmentation},
author = {Moenck, Keno and Schüppstuhl, Thorsten},
year = {2026},
url = {https://arxiv.org/abs/2608.30618},
doi = {https://doi.org/10.48550/arXiv.2608.30618}
}This codebase builds on and reuses components from several excellent open-source projects — in particular the data processing and backbone infrastructure of Pointcept, data preprocessing pipelines of Mask3D, the sparse UNet implementation of gorilla3d, the Volt point cloud backbone, and the SPFormer, MAFT, Relation3D, and OneFormer3D baselines. We thank the respective authors for open-sourcing their work.