coincheung / densecl

DenseCL + regionCL-D

DenseCL: Dense Contrastive Learning for Self-Supervised Visual Pre-Training

This is an unofficial PyTorch implementation of the DenseCL paper, with the help and suggestions from @WXinlong and @DerrickWang005.

Currently, regionCL-D is added, and pretrained checkpoints are uploaded.

Preparation

Install PyTorch and ImageNet dataset following the official PyTorch ImageNet training code.

This repo aims to be minimal modifications on that code. Check the modifications by:

diff main_densecl.py <(curl https://raw.githubusercontent.com/pytorch/examples/master/imagenet/main.py)
diff main_lincls.py <(curl https://raw.githubusercontent.com/pytorch/examples/master/imagenet/main.py)

Unsupervised Training & Linear Classification

This implementation only supports multi-gpu, DistributedDataParallel training, which is faster and simpler; single-gpu or DataParallel training is not supported.

This implementation only supports ResNet50/ResNet101, since we need to modify computing graph architecture and I only modified ResNet50/ResNet101.

To do unsupervised pre-training and linear-evaluation of a ResNet50/ResNet101 model on ImageNet in an 8-gpu machine, please refer to dist_train.sh for relevant starting script.

Since the paper says they use default mocov2 hyper-parameters, the above script uses same hyper-parameters as mocov2.

Note: for 4-gpu training, we recommend following the linear lr scaling recipe: --lr 0.015 --batch-size 128 with 4 gpus. We got similar results using this setting.

Models

Our pre-trained denseCL/RegionCL-D models can be downloaded as following:

	epochs	mlp	aug+	cos	IM top1	VOC AP50	model	md5
MoCov2 R50	200	✓	✓	✓	67.7	82.4	download	59fd9945
DenseCL R50	200	✓	✓	✓	63.8	82.7	download	7cfc894c
DenseCL R101	200	✓	✓	✓	65.4	83.5	download	006675e5
RegionCL-D R50	200	✓	✓	✓	67.5	83.3	download	8afad30e
RegionCL-D R101	200	✓	✓	✓	67.5	84.3	download	a1489ad4

Here IM is imagenet-1k dataset. We freeze pretrained weights and only fine tune the last classifier layer.

Please be aware that though DenseCL cannot match mocov2 in the filed of classification, it is superior to mocov2 in terms of object detection. More results of detection can be found here.

Transferring to Object Detection

For details, see ./detection.

License

This project is under the CC-BY-NC 4.0 license. See LICENSE for details.

regioncl, r50:
Acc@1 67.518 Acc@5 88.256
Acc@1 67.534 Acc@5 88.212
regioncl, r101:
Acc@1 67.504 Acc@5 88.212
Acc@1 67.470 Acc@5 88.104

densecl's Issues

an issue about shuffling BN

In your moco.py, I think 'dense_k' and 'feat_k' also need to be unshuffled back. Is that right?

About pretrained model of dense cl

i tried to download the pretrained model of dense cl but the link seems unavailable. Could you update the link?

training time

hi, thank you for sharing your codes.
I try to train densecl on imagenet with 4*2080ti, but it takes about 13 days to finish training, which is much longer time than yours.
Did you use imagenet in lmdb format so that training time is shorter?