Vector Neurons: A General Framework for SO(3)-Equivariant Networks

Created by Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas.

We introduce a general framework built on top of what we call Vector Neurons for creating SO(3) equivariant neural networks. Extending neurons from single scalars to 3D vectors, our vector neurons transport SO(3) actions to latent spaces and provide a framework for building equivariance in common neural operations including linear layers, non-linearities, pooling, and normalization.

Paper | Project | Short Video

🈶 Other Implementations

[Neural Implicit Reconstruction] Neural implicit reconstruction experiments in the paper.
[Pointcloud Networks] Separate training setups for VN-PointNet and VN-DGCNN.

Overview

vnn is the author's implementation of Vector Neuron Networks with PointNet and DGCNN backbones. The current version only supports input data without normals.

Data Preparation

Classification: Download ModelNet40 and save in data/modelnet40_normal_resampled/.
Part Segmentation: Download ShapeNet and save in data/shapenetcore_partanno_segmentation_benchmark_v0_normal/.

Usage

Classification on ModelNet40

Training

python train_cls.py --model vn_pointnet_cls --rot ROTATION --log_dir LOG_DIR

python train_cls.py --model vn_dgcnn_cls --rot ROTATION --log_dir LOG_DIR

Evaluation

python test_cls.py --model vn_pointnet_cls --rot ROTATION --log_dir LOG_DIR

python test_cls.py --model vn_dgcnn_cls --rot ROTATION --log_dir LOG_DIR

Here ROTATION should be chosen from aligned, z, so3. For instance, to train a VN-DGCNN on aligned shapes and test it on SO(3)-rotated shapes, run

python train_cls.py --model vn_dgcnn_cls --rot aligned --log_dir vn_dgcnn/aligned/

python test_cls.py --model vn_dgcnn_cls --rot so3 --log_dir vn_dgcnn/aligned/

Part Segmentation on ShapeNet

Training

python train_partseg.py --model vn_pointnet_partseg --rot ROTATION --log_dir LOG_DIR

python train_partseg.py --model vn_dgcnn_partseg --rot ROTATION --log_dir LOG_DIR

Evaluation

python test_partseg.py --model vn_pointnet_partseg --rot ROTATION --log_dir LOG_DIR

python test_partseg.py --model vn_dgcnn_partseg --rot ROTATION --log_dir LOG_DIR

For instance, to train a VN-DGCNN on aligned shapes and test it on SO(3)-rotated shapes, run

python train_partseg.py --model vn_dgcnn_partseg --rot aligned --log_dir vn_dgcnn/aligned/

python test_partseg.py --model vn_dgcnn_partseg --rot so3 --log_dir vn_dgcnn/aligned/

Citation

Please cite this paper if you want to use it in your work,

@article{deng2021vn,
  title={Vector Neurons: a general framework for SO(3)-equivariant networks},
  author={Deng, Congyue and Litany, Or and Duan, Yueqi and Poulenard, Adrien and Tagliasacchi, Andrea and Guibas, Leonidas},
  journal={arXiv preprint arXiv:2104.12229},
  year={2021}
}

License

MIT License

Acknowledgement

The structure of this codebase is borrowed from this pytorch implementataion of PointNet/PointNet++ and DGCNN as well as this implementation.

FlyingGiraffe / vnn