Skip to content

Latest commit

 

History

History
75 lines (64 loc) · 4.01 KB

README.md

File metadata and controls

75 lines (64 loc) · 4.01 KB
Raw

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.