Skip to content

Latest commit

 

History

History
57 lines (42 loc) · 2.48 KB

README.md

File metadata and controls

57 lines (42 loc) · 2.48 KB
Raw

Implement and reproduce results of the following papers:

Dependencies:

  • TensorFlow 1.14 or 1.15, built with XLA support
  • Tensorpack ≥ 0.10.1
  • Horovod ≥ 0.19 built with Gloo & NCCL support
  • TensorFlow zmq_ops
  • OpenCV
  • the taskset command (from the util-linux package)

Unsupervised Training:

To run MoCo pre-training on a machine with 8 GPUs, use:

horovodrun -np 8 --output-filename moco.log python main_moco.py --data /path/to/imagenet

Add --v2 to train MoCov2, which uses an extra MLP layer, extra augmentations, and cosine LR schedule.

Linear Classification:

To train a linear classifier using the pre-trained features, run:

./main_lincls.py --load /path/to/pretrained/checkpoint --data /path/to/imagenet

KNN Evaluation:

Instead of Linear Classification, a cheap but rough evaluation is to perform a feature-space kNN using the training set:

horovodrun -np 8 ./eval_knn.py --load /path/to/checkpoint --data /path/to/imagenet --top-k 200

Results:

Training was done in a machine with 8 V100s, >200GB RAM and 80 CPUs.

Following results are obtained after 200 epochs of pre-training (~53h) and 100 epochs of linear classifier tuning (~8h). KNN evaluation takes 10min per checkpoint.

linear cls.
accuracy		 download
(pretrained only)		 tensorboard
MoCo v1 60.9% ⬇️ N/A
MoCo v2 67.7% ⬇️ pretrain; finetune

Notes:

  • Horovod with Gloo is recommended. Horovod with MPI is not tested and may crash due to how we use forking.
  • If using TensorFlow without XLA support, you can modify main_*.py to replace xla.compile by a naive forward.
  • Official PyTorch code is at facebookresearch/moco.