Learning SO(3) Equivariant Representations with Spherical CNNs

Abstract

We address the problem of 3D rotation equivariance in convolutional neural networks. 3D rotations have been a challenging nuisance in 3D classification tasks requiring higher capacity and extended data augmentation in order to tackle it. We model 3D data with multi-valued spherical functions and we propose a novel spherical convolutional network that implements exact convolutions on the sphere by realizing them in the spherical harmonic domain. Resulting filters have local symmetry and are localized by enforcing smooth spectra. We apply a novel pooling on the spectral domain and our operations are independent of the underlying spherical resolution throughout the network. We show that networks with much lower capacity and without requiring data augmentation can exhibit performance comparable to the state of the art in standard retrieval and classification benchmarks.

Demo

This repository contains a demo, where we train and test the model on the SO(3)-rotated ModelNet40 dataset.

Check the requirements in requirements.txt. Our codebase has been tested on TensorFlow 1.6 but the dependency is commented out to silence GitHub’s security warnings.

Download the dataset here (7.0Gb).

The following commands should

• uncompress the downloaded dataset from /tmp to ~/data/m40-so3-64,
• clone this repo,
• create a virtualenv,
• install the requirements,
• train and test the model.

mkdir -p ~/data/m40-so3-64 && tar -xvf /tmp/m40_tf_csph_3d_48aug_64.tar.gz -C ~/data/m40-so3-64

git clone https://github.com/daniilidis-group/spherical-cnn.git
cd spherical-cnn

virtualenv -p python3 env
source env/bin/activate
pip install -r requirements.txt
TF_CPP_MIN_LOG_LEVEL=1 python3 scripts/train.py \
                               @params/model-64.txt \
                               @params/m40-64.txt \
                               @params/training.txt \
                               --dset_dir ~/data/m40-so3-64 \
                               --logdir /tmp/m40-so3 \
                               --run_id m40-so3

Sample output:

Running on m40-so3
Namespace(...)
Loading dataset from from_cached_tfrecords...
Loading model two_branch. Logdir /tmp/m40-so3
Start training...
epoch=1; lr=0.0010 train: 0.2286, valid: 0.2286
epoch=2; lr=0.0010 train: 0.4335, valid: 0.4335
epoch=3; lr=0.0010 train: 0.5975, valid: 0.5975
(...)
epoch=46; lr=0.0000 train: 0.9630, valid: 0.9630
epoch=47; lr=0.0000 train: 0.9640, valid: 0.9640
epoch=48; lr=0.0000 train: 0.9724, valid: 0.9724
Start testing...
| model   |  train |    val |   test | train time |
| m40-so3 | 0.9724 | 0.9724 | 0.8679 |      88.31 |

This reproduces the SO(3)/SO(3) result in table 1 (86.9%) of the paper. It runs in about 75min on a Nvidia GeForce GTX 1080 Ti, and requires ~2.2 Gb of GPU memory.

Change --dset_dir if you saved the dataset elsewhere.

Call pytest to run the unit tests.

References

Esteves, C., Allen-Blanchette, C., Makadia, A., & Daniilidis, K. Learning SO(3) Equivariant Representations with Spherical CNNs. European Conference on Computer Vision, ECCV 2018 (oral). http://arxiv.org/abs/1711.06721

@article{esteves17_learn_so_equiv_repres_with_spher_cnns,
  author = {Esteves, Carlos and Allen-Blanchette, Christine and Makadia, Ameesh and Daniilidis, Kostas},
  title = {Learning SO(3) Equivariant Representations With Spherical Cnns},
  journal = {CoRR},
  year = {2017},
  url = {http://arxiv.org/abs/1711.06721},
  archivePrefix = {arXiv},
  eprint = {1711.06721},
  primaryClass = {cs.CV},
}

Authors

Carlos Esteves [1], Christine Allen-Blanchette [1], Ameesh Makadia [2], Kostas Daniilidis [1]

[1] GRASP Laboratory, University of Pennsylvania

[2] Google