Equivariant Vision Models

This repository offers implementations of common vision models that are equivariant to rotations and reflections in the image plane. For most models, we also offer weights achieved by pretraining the network on ImageNet-1k.

List of Models

Name	Num. Params	Acc@1	Acc@5	Checkpoint
d1resnet18	11.5M	0.709	0.901	ckpt
c4resnet18	11.7M	0.734	0.915	ckpt
d4resnet18	11.7M	0.737	0.916	ckpt
c8resnet18	11.7M	0.738	0.914	ckpt
d8resnet18	11.9M	0.736	0.915	ckpt
d1resnet50	25.7M	0.769	0.935	ckpt
c4resnet50	24.7M	0.785	0.943	ckpt
d4resnet50	24.8M	0.789	0.946	ckpt
c8resnet50	24.8M	0.787	0.945	ckpt
d1resnet101	44.7M	0.785	0.937	ckpt
c4resnet101	43.4M	0.801	0.952	ckpt
d4resnet101	43.9M	0.804	0.953	ckpt
c8resnet101	43.9M

If you would like to suggest a new model or contribute an implementation, open an issue or PR. We may be able to train the model for you.

Installation

After downloading the repository, you can install the necessary packages for using the models with:

pip install -r requirements.txt

Alternatively, you can install the package with PyPI:

pip install git+https://github.com/dmklee/equivision

Usage

Below is a code snippet that illustrates how to load a pretrained model and run inference on it. The model is composed of equivariant modules so it will maintain equivariance even if trained further.

import torch
from equivision import models

model = models.d1resnet18(pretrained=True)

dummy_img = torch.zeros((1, 3, 224, 224))

class_logits = model.forward(dummy_img)
equiv_fmap = model.forward_features(dummy_img)

Note that the model takes as input a torch.Tensor, not an escnn.nn.GeometricTensor. The result of model.forward_features is a GeometricTensor with regular features.

Training

The following command will create a virtual environment with all necessary packages for training.

./setup.sh

To train a model, use the train.py script. The following is the command used to replicate the model weights stored in this repository:

python train.py --num_workers=6 --devices=4 --data_dir=/path/to/imagenet \
	--seed=10 --accumulate_grad_batches=2 --batch_size=32 --precision=16-mixed \
	--model=<model-name>

Acknowledgements

The training procedure we used was developed by PyTorch (see here). The models are implemented using escnn and inspired by the wide resnet example here. Please consider citing their work if you use these models.