Train RotNet model

VISSL reproduces the self-supervised approach Unsupervised Representation Learning by Predicting Image Rotations proposed by Spyros Gidaris, Praveer Singh, Nikos Komodakis in https://arxiv.org/abs/1803.07728.

How to train RotNet model

VISSL provides yaml configuration file containing the exact hyperparam settings to reproduce the model. VISSL implements all the components including data augmentations, collators etc required for this approach.

To train ResNet-50 model on 8-gpus on ImageNet-1K dataset using 4 rotation angles:

python tools/run_distributed_engines.py config=pretrain/rotnet/rotnet_8gpu_resnet

Training different model architecture

VISSL supports many backbone architectures including AlexNet, ResNe(X)ts. Some examples below:

  • Train AlexNet model

python tools/run_distributed_engines.py config=pretrain/rotnet/rotnet_8gpu_resnet \
    config.MODEL.TRUNK.NAME=alexnet_rotnet
  • Train ResNet-101:

python tools/run_distributed_engines.py config=pretrain/rotnet/rotnet_8gpu_resnet \
    config.MODEL.TRUNK.NAME=resnet config.MODEL.TRUNK.TRUNK_PARAMS.RESNETS.DEPTH=101

Vary the number of gpus

VISSL makes it extremely easy to vary the number of gpus to be used in training. For example: to train the RotNet model on 4 machines (32gpus) or 1gpu, the changes required are:

  • Training on 1-gpu:

python tools/run_distributed_engines.py config=pretrain/rotnet/rotnet_8gpu_resnet \
    config.DISTRIBUTED.NUM_PROC_PER_NODE=1
  • Training on 4 machines i.e. 32-gpu:

python tools/run_distributed_engines.py config=pretrain/rotnet/rotnet_8gpu_resnet \
    config.DISTRIBUTED.NUM_PROC_PER_NODE=8 config.DISTRIBUTED.NUM_NODES=4

Note

Please adjust the learning rate following ImageNet in 1-Hour if you change the number of gpus.

Pre-trained models

See VISSL Model Zoo for the PyTorch pre-trained models with VISSL using RotNet approach and the benchmarks.

Citation

@misc{gidaris2018unsupervised,
    title={Unsupervised Representation Learning by Predicting Image Rotations},
    author={Spyros Gidaris and Praveer Singh and Nikos Komodakis},
    year={2018},
    eprint={1803.07728},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}