Unofficial implementation for MoCo: Momentum Contrast for Unsupervised Visual Representation Learning

This repository carefully implemented important details such as ShuffleBN and distributed Queue mentioned in the paper to reproduce the reported results.

The following enverionments is tested:

• Anaconda with python >= 3.6
• pytorch>=1.3, torchvision, cuda=9.2
• tensorboard_logger: pip install tensorboard_logger

• The pre-training stage:

  exp_name=MoCo/ddp/cifar10_test
  python -m torch.distributed.launch --nproc_per_node=1 \
      train_cifar_no_dist.py \
      --batch-size 32 \
      --exp-name ${exp_name} \
      --learning-rate 0.01 \
      --nce-k 4096 \
      --nce-t 0.3 \
      --epochs 500 \
      --weight-decay 5e-4
  

• The linear evaluation stage:

  exp_name=MoCo/ddp/cifar10_test
  python -m torch.distributed.launch --nproc_per_node=1 \
      eval_cifar_no_dist.py \
      --exp-name ${exp_name} \
      --model-path ./output/cifar10/${exp_name}/models/current.pth \
      --batch-size 128 \
      --learning-rate 3 \
      --epochs 20 \
      --cosine
  

• The pre-training stage:

  exp_name=MoCo/ddp/8-gpu_bs-256_shuffle_bn
  python -m torch.distributed.launch --nproc_per_node=8 \
      train.py \
      --batch-size 32 \
      --exp-name ${exp_name} \
      --data-root /root/directory/of/datasets
  

  The checkpoints and tensorboard log will be saved in ./output/imagenet/${exp_name}. Run python train.py --help for more help.

• The linear evaluation stage:

  exp_name=MoCo/ddp/8-gpu_bs-256_shuffle_bn
  python -m torch.distributed.launch --nproc_per_node=4 \
      eval.py \
      --exp-name ${exp_name} \
      --model-path ./output/imagenet/${exp_name}/models/current.pth \
      --batch-size 64 \
      --data-root /root/directory/of/datasets
  

  The checkpoints and tensorboard log will be saved in ./output/imagenet/${exp_name}. Run python eval.py --help for more help.

Pre-trained model checkpoint and tensorboard log for K = 16384 and 65536 can be downloaded from OneDrive.

A lot of codes is borrowed from CMC and lemniscate.