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 this paper.

How to train RotNet model¶

VISSL provides a yaml configuration file containing the exact hyperparameter 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.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}
}