Group-equivariant CNN (G-CNN)

This project is a pytorch implementation of Group-equivariant Convolutional Neural Networks by Cohen and Welling¹. We create custom G-convolutional layers and G-pooling layers for various noncommutative groups of interest, namely:

p4 (translations and rotations of 90 degrees)
p4m (translations and rotations of 90 degrees + reflections)

There are a lot of similarities with Cohen and Welling's original implementation here. The main differences are that Cohen and Welling used Tensorflow and CuPy for GPU acceleration, whereas I used Pytorch and Libtorch for extensions with CUDA kernels.

Requirements

python >= 3.7
numpy
pytorch >= 1.12.1
CUDA 11.3 (CUDA 11.6 doesn't currently work with the GPU kernels)
libtorch
scikit-image
scikit-learn
Microsoft Visual Studio (for Windows users)

Usage

Before running the code it is highly recommended to install the C++/CUDA optimization modules. (Otherwise it may be very slow!)

In the command line, use the following commands
py setup_cpp.py install (C++ extensions)
py setup_cuda.py install (CUDA extensions)

These will install the corresponding C++/CUDA functions instead of the native pytorch functions used for forward() and backward() functions in the GCNN models.

The training code is in train.py. There are two models available currently:

• P4CNN: used to train rotated MNIST
• P4AllCNN: used to train (rotated) CIFAR10

To start training, use the command py train.py --model=p4cnn or py train.py --model=p4allcnn

Task List

• Optimization using C++/CUDA
• Further GPU optimization with shared memory
• Dataloader with rotations and reflections support
• Unit tests
• Training scripts
• Benchmark results
• Visualizations

---

Footnotes

1. Taco S. Cohen, Max Welling Group Equivariant Convolutional Networks (2016) ↩