This project is a pytorch implementation of Group-equivariant Convolutional Neural Networks by Cohen and Welling1.
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.
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)
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
- Optimization using C++/CUDA
- Further GPU optimization with shared memory
- Dataloader with rotations and reflections support
- Unit tests
- Training scripts
- Benchmark results
- Visualizations
Footnotes
-
Taco S. Cohen, Max Welling Group Equivariant Convolutional Networks (2016) ↩