This repository is the official PyTorch implementation of
VectorAdam for Rotation Equivariant Geometry Optimization
Selena Ling, Nicholas Sharp, Alec Jacobson
NeurIPS 2022
To use our VectorAdam implementation, you just need to have PyTorch installed in your environment.
The demo script are tested with PyTorch=1.11 and matplotlib=3.5.1. We also provide the environment file vectoradam.yml, which can be used to create a conda environment as in
conda env create -f vectoradam.yml -n [env-name]
Note that this is tested on Ubuntu 18.04 only.
To use VectorAdam in your project,
optimizer = VectorAdam(
[{'params': X, 'axis': -1},
{'params': Y, 'axis': 1],
lr=lr, betas=betas, eps=eps))
The above example will apply VectorAdam's vector-wise operations to X along the last axis and Y along the 1st axis, with specified learning rate, betas and epsilon hyperparameters.
We provide a demo with laplacian2d_demo.ipynb that reproduces the 2D results we have in Figure 4.