This is the code of our paper "Semantic-Aware Auto-Encoders for Self-Supervised Representation Learning".
Guangrun Wang, Yansong Tang, Liang Lin, and Philip H.S. Torr.
A highly recommended project: Traditional Classification Neural Networks are Good Generators: They are Competitive with DDPMs and GANs
An example of SSL training script on ImageNet. More and larger GPUs are better. Please refer to DINO to use multi-node training).
python -m torch.distributed.launch --nproc_per_node=8 main_saae.py --arch vit_base --data_path xxxxxxxxx/ILSVRC2012/train --output_dir xxxxxxx/name_of_output_dir/ --epochs 400 --reference_crops_number 1 --batch_size_per_gpu 120 --use_fp16 true --num_workers 4 --weight_decay 0.1 --weight_decay_end 0.1 --use_bn_in_head true --lr 2.4e-4 --global_crops_scale 0.14 1.
An example of linear evaluation script on ImageNet:
python3 -m torch.distributed.launch --nproc_per_node=8 eval_linear.py --arch vit_base --checkpoint_key student --num_workers 4 --data_path xxxxxxxx/ILSVRC2012 --pretrained_weights xxxxxxxxx/name_of_output_dir/checkpoint.pth --output_dir xxxxxxxxxxxx/name_of_eval_output_dir/ --lr 0.01
This repository is released under the Apache 2.0 license as found in the LICENSE file.
This repo has been tested in the following environment. More precisely, this repo is a modification on the DINO. Installation and preparation follow that repo. Please acknowledge the great work of the team of DINO. Thanks for their outstanding contributions.
Pytorch1.9
Solving Inefficiency of Self-supervised Representation Learning