cyx669521 / self-supervised-learning

Fast, Simple and Easy to use SOTA Self-Supervised Image Classification Models in PyTorch

Self Supervised Image Classification

Introduction

Read a blog post from FAIR >> Self-supervised learning: The dark matter of intelligence.

Models

• EsViT
• XCiT
• DINO
• MoCov3

Model Zoo

Method	Model	ImageNet Top1 Acc (Linear)	ImageNet Top1 Acc (k-NN)	Params (M)	Weights
EsViT	Swin-B/W=14	81.3	79.3	87	N/A
EsViT	Swin-S/W=14	80.8	79.1	49	N/A
EsViT	Swin-T/W=14	78.7	77.0	28	N/A
DINO	XCiT-M24/8	80.3	77.9	84	model/checkpoint
DINO	XCiT-S12/8	79.2	77.1	26	model/checkpoint
DINO	ViT-B/8	80.1	77.4	85	model/checkpoint
DINO	ViT-S/8	79.7	78.3	21	model/checkpoint
MoCov3	ViT-B	76.7	-	-	N/A
MoCov3	ViT-S	73.2	-	-	N/A

Configuration

Create a configuration file in configs. Sample configuration for ImageNet dataset with DINO can be found here. Then edit the fields you think if it is needed. This configuration file is needed for all of training, evaluation and prediction scripts.

Training

Single GPU

$ python tools/train.py --cfg configs/CONFIG_FILE_NAME.yaml

Multiple GPUs

Traing with 2 GPUs:

$ python -m torch.distributed.launch --nproc_per_node=2 --use_env tools/train.py --cfg configs/CONFIG_FILE_NAME.yaml

Evaluation

Make sure to set MODEL_PATH of the configuration file to your trained model directory.

Linear Classification

This will train a supervised linear classifier on top of trained weights and evaluate the result.

$ python -m torch.distributed.launch --nproc_per_node=2 --use_env tools/val_linear.py --cfg configs/CONFIG_FILE_NAME.yaml

k-NN Classification

$ python -m torch.distributed.launch --nproc_per_node=1 --use_env tools/val_knn.py --cfg configs/CONFIG_FILE_NAME.yaml

Attention Visualization

Make sure to set MODEL_PATH of the configuration file to model's weights.

$ python tools/visualize_attention.py --cfg configs/CONFIG_FILE_NAME.yaml