CHOIR: CHaracteristic Orientation Predictor with Invariant Residual Learning

Stable and Consistent Prediction of 3D Characteristic Orientation via Invariant Residual Learning
Seungwook Kim^1*, Chunghyun Park^1*, Jaesik Park², and Minsu Cho¹ (*equal contribution)
¹POSTECH and ²Seoul National University
ICML 2023, Honolulu.

Installation

curl -LsSf https://astral.sh/uv/install.sh | sh
uv sync

Note

CUDA is required for torch-cluster (compiled from source during uv sync). Tested on Python 3.12, PyTorch 2.11, CUDA 12.8, 8 NVIDIA A6000 GPUs.

Data

ShapeNet point clouds are automatically downloaded and prepared on first run. To prepare manually:

uv run python src/data/prepare.py

This downloads ShapeNet from AtlasNetV2, converts PLY to H5, and generates the stability evaluation set.

Training

# Single-class (8 GPU DDP)
uv run python src/train.py experiment=airplane
uv run python src/train.py experiment=car
uv run python src/train.py experiment=chair
uv run python src/train.py experiment=table

# Multi-class (airplane, car, chair, table)
uv run python src/train.py experiment=multi

# Override any config
uv run python src/train.py experiment=airplane trainer.devices=1 data.batch_size=16

Checkpoints and logs are saved to logs/.

Evaluation

uv run python src/eval.py ckpt_path=logs/.../best.ckpt

Acknowledgments

This project builds upon Vector Neurons for SO(3)-equivariant layers and Canonical Capsules for ShapeNet data processing. The codebase structure follows lightning-hydra-template.

Citation

If you find our work useful, please consider citing:

@inproceedings{kim2023choir,
  title={Stable and Consistent Prediction of 3D Characteristic Orientation via Invariant Residual Learning},
  author={Kim, Seungwook and Park, Chunghyun and Park, Jaesik and Cho, Minsu},
  booktitle={Proceedings of the International Conference on Machine Learning (ICML)},
  year={2023}
}