|
๐ Paper arXiv:2506.21724 |
๐ Conference NeurIPS 2025 (accepted) |
๐ง๐ปโ๐ผ Authors Remco F. Leijenaar Hamidreza Kasaei |
Overview of the Asymmetric Dual Self-Distillation (AsymDSD) framework.
If you find this repository useful, please cite our paper:
@article{leijenaar2025asymmetric,
title={Asymmetric Dual Self-Distillation for 3D Self-Supervised Representation Learning},
author={Leijenaar, Remco F and Kasaei, Hamidreza},
journal={Advances in Neural Information Processing Systems},
year={2025}
}Ensure your system provides the following:
- Python 3.11
- GCC/G++ 6โ13.3
- C++ build tools (
g++,make; e.g.build-essentialon Ubuntu) - Python development headers (
python3-dev) git
- CUDA compiler (
nvcc, from a CUDA Toolkit installation or Conda CUDA package)
โ ๏ธ Other versions may work, but only the above configuration has been tested.
Set up the environment using uv:
uv venv
source .venv/bin/activate
uv pip install -r requirements.txt --index-strategy unsafe-best-match
uv pip install git+https://github.com/facebookresearch/pytorch3d.git@stable --no-build-isolation
uv pip install -e .If you prefer Conda, we recommend using Mamba via Miniforge:
mamba env create -f conda.yaml
conda activate asymdsdpip install -e .PyTorch3D is built from source (it is not provided as a prebuilt wheel).
If you want CUDA support, you need a CUDA compiler (nvcc) available either via a system CUDA Toolkit install or via a Conda-provided cuda-nvcc/toolkit package.
If you run into issues during setup, try (exclude uv if using Conda/Mamba):
export MAX_JOBS=4
uv pip install git+https://github.com/facebookresearch/pytorch3d.git@stable --no-build-isolation๐ข Lowering
MAX_JOBSreduces peak memory usage (slower but useful for systems with limited memory).
With the provided configurations, dataset files should be placed inside the data folder. Do not extract the archive filesโthe code reads directly from the compressed archives.
- Request access via ShapeNet/ShapeNetCore-archive on Hugging Face.
- After approval, download
ShapeNetCore.v2.zipand place it in thedatafolder.
wget -P data http://modelnet.cs.princeton.edu/ModelNet40.zip๐ For more information, visit the ModelNet40 project page.
- Visit the ScanObjectNN website and agree to the terms of use.
- Download the dataset
h5_files.zipand place it in thedata/ScanObjectNNdirectory.
mkdir -p data/ScanObjectNN
# Replace the placeholder below with the actual download link after gaining access
wget -P data/ScanObjectNN <DOWNLOAD_LINK>Download the ShapeNetPart archive into data/:
cd data
gdown https://drive.google.com/uc?id=1W3SEE-dY1sxvlECcOwWSDYemwHEUbJISThis should create:
data/shapenetcore_partanno_segmentation_benchmark_v0_normal.tar
- Download the dataset from ModelNet40 Few-Shot, by selecting all files and downloading them as a zip file
ModelNetFewshot.zip. - Place the zip file in the data folder:
data.
โน๏ธ Training on the Mixture dataset requires additional datasets beyond the ones listed in this section. See Section 5 for the extra preparation steps.
To start pretraining the small version of AsymDSD on ShapeNetCore, run:
sh shell_scripts/sh/train_ssrl.sh๐งญ You may be prompted to log in to Weights & Biases (wandb).
The first time you run this, it will compile and preprocess the datasets. This process may take a while, but all data is cached under the data directory, making subsequent runs much faster.
To train with specific modes, use the corresponding configuration files:
-
MPM mode:
sh shell_scripts/sh/train_ssrl.sh --model configs/ssrl/variants/model/ssrl_model_mask.yaml
-
CLS mode:
sh shell_scripts/sh/train_ssrl.sh --model configs/ssrl/variants/model/ssrl_model_cls.yaml
Train the base-sized model on the Mixture dataset:
sh shell_scripts/sh/train_ssrl.sh --model configs/ssrl/variants/model/ssrl_model_base.yaml --data configs/data/Mixture-U.yaml๐ก To accelerate pre-training, you can disable evaluation callbacks by editing the trainer config file, or skip all callbacks by passing
--trainer.callbacks null
To evaluate the model on object recognition tasks, use the following command:
python shell_scripts/py/train_neural_classifier_all.py --runs <num_eval_runs> --model.encoder_ckpt_path <path_to_model>For the MPM pre-trained version without a CLS-token, you can add:
--model configs/classification/variants/model/classification_model_mask.yamlFor few-shot evaluation on ModelNet40:
python shell_scripts/py/train_neural_classifier_all.py --model.encoder_ckpt_path <path_to_model>To run semantic segmentation fine-tuning/evaluation, run:
sh shell_scripts/sh/train_semseg.sh --model.encoder_ckpt_path <ckpt>If you prepared Objaverse-v2 with LVIS annotations (see Section 5), you can run LVIS few-shot evaluation with:
python shell_scripts/py/Objaverse_fewshot_evals.py --model.encoder_ckpt_path <path_to_model>Tip: To evaluate a base-sized encoder, add
--model configs/classification/variants/model/classification_model_base.yaml.
๐ You can find logged results on Weights and Biases. A link to the run is provided in the script output.
We provide AsymDSD pre-trained model checkpoints on Hugging Face:
- Hugging Face repo: https://huggingface.co/remcofl/AsymDSD
- Suggested local folder:
checkpoints/
| Checkpoint file (HF) | Size | Pre-training data | Variant | Params (M) |
|---|---|---|---|---|
AsymDSD-S_ShapeNet.ckpt |
230 MB | ShapeNetCore | AsymDSD-S | 21.8 |
AsymSD-S-CLS_ShapeNet.ckpt |
197 MB | ShapeNetCore | AsymSD-S-CLS | 21.8 |
AsymSD-S-MPM_ShapeNet.ckpt |
208 MB | ShapeNetCore | AsymSD-S-MPM | 21.8 |
AsymDSD-B_Mixture_B.ckpt |
827 MB | Mixture | AsymDSD-B | 91.8 |
The environment includes huggingface_hub, which provides the hf CLI.
Download a single checkpoint into checkpoints/:
hf download remcofl/AsymDSD <model.ckpt> --local-dir checkpoints/If you want to download multiple files, repeat the command with a different filename from the table above.
Once downloaded, you can point evaluation/fine-tuning scripts to the checkpoint, e.g.:
python shell_scripts/py/train_neural_classifier_all.py --runs <num_eval_runs> --model.encoder_ckpt_path checkpoints/AsymDSD-S_ShapeNet.ckptThe Mixture dataset is configured in configs/data/Mixture-U.yaml.
For most sources, the dataset cache is built automatically the first time you run training (if the raw data is present).
If you want to pre-build caches (or if a dataset is not auto-prepared), use the provided CLI wrapper:
sh shell_scripts/sh/prepare_data_zarr.sh <dataset_name|config_path>In addition to the datasets listed in Section 2, Mixture uses the following sources.
- Download the collection using the script shipped in this repository:
mkdir -p data/ScannedObjects
cd data/ScannedObjects
python ../../asymdsd/data/datasets_/ScannedObjects/download_collection.py \
-o "GoogleResearch" \
-c "Scanned Objects by Google Research"
cd ../..- Optional: Prepare the dataset cache:
sh shell_scripts/sh/prepare_data_zarr.sh ScannedObjectsConfig: configs/data/prepare_data_zarr/ScannedObjects.yaml
- Go to: https://openxlab.org.cn/datasets/OpenXDLab/OmniObject3D-New/tree/main/raw/point_clouds/hdf5_files
- Create an OpenXLab account and log in if required.
- Download the
16384point cloud files. - Create
data/OmniObject3D.zipcontaining the downloaded files (or rename your downloaded archive toOmniObject3D.zip). - Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/OmniObject3D.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh OmniObject3D- Go to: https://tianchi.aliyun.com/dataset/98063
- Create an account, log in, and click Apply for dataset to request access.
- Place the model-part archives in
data/3D-Future/:
data/3D-Future/3D-FUTURE-model-part1.zipdata/3D-Future/3D-FUTURE-model-part2.zipdata/3D-Future/3D-FUTURE-model-part3.zipdata/3D-Future/3D-FUTURE-model-part4.zip
- Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/3D-FUTURE.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh 3D-FUTURE- Go to: https://github.com/rehg-lab/lowshot-shapebias/tree/main/toys4k
- In the Downloading Toys4K section, click the provided link and fill in the form to request access.
- Place the archive at:
data/toys4k_obj_files.zip - Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/Toys4K.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh Toys4K- Request access by filling in this form: https://docs.google.com/forms/d/e/1FAIpQLScDimvNMCGhy_rmBA2gHfDu3naktRm6A8BPwAWWDv-Uhm6Shw/viewform?c=0&w=1
- Download
Stanford3dDataset_v1.2.zipand place it at:data/Stanford3dDataset_v1.2.zip - Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/S3DIS_objects.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh S3DIS_objects- Download the required files into
data/SUNRGBD/:
mkdir -p data/SUNRGBD
cd data/SUNRGBD
wget https://rgbd.cs.princeton.edu/data/SUNRGBD.zip
wget https://rgbd.cs.princeton.edu/data/SUNRGBDtoolbox.zip
wget https://rgbd.cs.princeton.edu/data/SUNRGBDMeta3DBB_v2.mat
cd ../..- Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/SUNRGBD.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh SUNRGBD- Download the required files into
data/ABO/:
mkdir -p data/ABO
cd data/ABO
wget https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-3dmodels.tar
wget https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-listings.tar
cd ../..- Optional: Prepare the dataset cache (config:
configs/data/prepare_data_zarr/ABO.yaml):
sh shell_scripts/sh/prepare_data_zarr.sh ABOMixture expects a prepared data/Objaverse.zarr. Due to the long runtime, this is not prepared automatically as part of pre-training.
You can prepare Objaverse in two ways:
- Objaverse-v2 (recommended): more reliable and faster.
This option also prepares LVIS-based splits/labels used for LVIS few-shot evaluation.
mkdir -p data/ObjaverseV2
sh shell_scripts/sh/prepare_data_zarr.sh Objaverse_v2โน๏ธ No files need to be manually downloaded into
data/ObjaverseV2for this option. The script will handle downloading and preparing of the raw files.
Config: configs/data/prepare_data_zarr/Objaverse_v2.yaml
- Objaverse-v1: use this option if you want to stay closer to replicating the Mixture dataset used in the paper.
mkdir -p data/Objaverse
sh shell_scripts/sh/prepare_data_zarr.sh ObjaverseConfig: configs/data/prepare_data_zarr/Objaverse.yaml
๐ก Re-running the command resumes and processes remaining objects.
Planned future releases:
- Pre-trained Models: Checkpoints for both small and base versions of AsymDSD, including AsymDSD-CLS and AsymDSD-MPM.
- Additional Datasets: Dataset preparation modules including Mixture and Objaverse.
- Training Scripts: Full training configurations for larger model variants and part segmentation on ShapeNet-Part.
