Skip to content

CoinCheung/DenseCL

Repository files navigation

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