SoCo

[NeurIPS 2021 Spotlight] Aligning Pretraining for Detection via Object-Level Contrastive Learning

By Fangyun Wei*, Yue Gao*, Zhirong Wu, Han Hu, Stephen Lin.

* Equal contribution.

Introduction

Image-level contrastive representation learning has proven to be highly effective as a generic model for transfer learning. Such generality for transfer learning, however, sacrifices specificity if we are interested in a certain downstream task. We argue that this could be sub-optimal and thus advocate a design principle which encourages alignment between the self-supervised pretext task and the downstream task. In this paper, we follow this principle with a pretraining method specifically designed for the task of object detection. We attain alignment in the following three aspects:

1. object-level representations are introduced via selective search bounding boxes as object proposals;
2. the pretraining network architecture incorporates the same dedicated modules used in the detection pipeline (e.g. FPN);
3. the pretraining is equipped with object detection properties such as object-level translation invariance and scale invariance. Our method, called Selective Object COntrastive learning (SoCo), achieves state-of-the-art results for transfer performance on COCO detection using a Mask R-CNN framework.

Architecture

Main results

The pretrained models and finetuned models with their logs are available on Google Drive and Baidu Pan (code: 4662)

The following links are relative paths of the share folder.

SoCo pre-trained models

Model	Arch	Epochs	Scripts	Pretrained Model (relative path)
SoCo	ResNet50-C4	100	SoCo_C4_100ep	log (pretrain/SoCo_C4_100ep/log.txt) raw model (pretrain/SoCo_C4_100ep/ckpt_epoch_100.pth) converted d2 model (pretrain/SoCo_C4_100ep/current_detectron2_C4.pkl)
SoCo	ResNet50-C4	400	SoCo_C4_400ep	log (pretrain/SoCo_C4_400ep/log.txt) raw model (pretrain/SoCo_C4_400ep/ckpt_epoch_400.pth) converted d2 model (pretrain/SoCo_C4_400ep/current_detectron2_C4.pkl)
SoCo	ResNet50-FPN	100	SoCo_FPN_100ep	log (pretrain/SoCo_FPN_100ep/log.txt) raw model (pretrain/SoCo_FPN_100ep/ckpt_epoch_100.pth) converted d2 model (pretrain/SoCo_FPN_100ep/current_detectron2_Head.pkl)
SoCo	ResNet50-FPN	400	SoCo_FPN_400ep	log (pretrain/SoCo_FPN_400ep/log.txt) raw model (pretrain/SoCo_FPN_400ep/ckpt_epoch_400.pth) converted d2 model (pretrain/SoCo_FPN_400ep/current_detectron2_Head.pkl)
SoCo*	ResNet50-FPN	400	SoCo_FPN_Star_400ep	log (pretrain/SoCo_FPN_Star_400ep/log.txt) raw model (pretrain/SoCo_FPN_Star_400ep/ckpt_epoch_400.pth) converted d2 model (pretrain/SoCo_FPN_Star_400ep/current_detectron2_Head.pkl)

Results on LVIS with MaskRCNN R50-FPN

Methods	Epoch	APbb	APbb50	APbb75	APmk	APmk50	APmk75	config	Detectron2 trained (relative path)
Supervised	90	20.4	32.9	21.7	19.4	30.6	20.5	--	--
SoCo*	400	26.3	41.2	27.8	25.0	38.5	26.8	config	log (finetune/mask_rcnn_lvis_SoCo_FPN_Star_400ep_1x/log.txt) model (finetune/mask_rcnn_lvis_SoCo_FPN_Star_400ep_1x/model_final.pth)

Results on COCO with MaskRCNN R50-FPN

Methods	Epoch	APbb	APbb50	APbb75	APmk	APmk50	APmk75	config	Detectron2 trained (relative path)
Scratch	-	31.0	49.5	33.2	28.5	46.8	30.4	--	--
Supervised	90	38.9	59.6	42.7	35.4	56.5	38.1	--	--
SoCo	100	42.3	62.5	46.5	37.6	59.1	40.5	config	log (finetune/mask_rcnn_coco_SoCo_FPN_100ep_1x/log.txt) model (finetune/mask_rcnn_coco_SoCo_FPN_100ep_1x/model_final.pth)
SoCo	400	43.0	63.3	47.1	38.2	60.2	41.0	config	log (finetune/mask_rcnn_coco_SoCo_FPN_400ep_1x/log.txt) model (finetune/mask_rcnn_coco_SoCo_FPN_400ep_1x/model_final.pth)
SoCo*	400	43.2	63.5	47.4	38.4	60.2	41.4	config	log (finetune/mask_rcnn_coco_SoCo_FPN_Star_400ep_1x/log.txt) model (finetune/mask_rcnn_coco_SoCo_FPN_Star_400ep_1x/model_final.pth)

Results on COCO with MaskRCNN R50-C4

Methods	Epoch	APbb	APbb50	APbb75	APmk	APmk50	APmk75	config	Detectron2 trained (relative path)
Scratch	-	26.4	44.0	27.8	29.3	46.9	30.8	--	--
Supervised	90	38.2	58.2	41.2	33.3	54.7	35.2	--	--
SoCo	100	40.4	60.4	43.7	34.9	56.8	37.0	config	log (finetune/mask_rcnn_coco_SoCo_C4_100ep_1x/log.txt) model (finetune/mask_rcnn_coco_SoCo_C4_100ep_1x/model_final.pth)
SoCo	400	40.9	60.9	44.3	35.3	57.5	37.3	config	log (finetune/mask_rcnn_coco_SoCo_C4_400ep_1x/log.txt) model (finetune/mask_rcnn_coco_SoCo_C4_400ep_1x/model_final.pth)

Get started

Requirements

The Dockerfile is included, please refer to it.

Prepare data with Selective Search

1. Generate Selective Search proposals

   python selective_search/generate_imagenet_ss_proposals.py

2. Filter out invalid proposals with filter strategy

   python selective_search/filter_ss_proposals_json.py

3. Post preprocessing for images of no proposals

   python selective_search/filter_ss_proposals_json_post_no_prop.py

Pretrain with SoCo

Use SoCo FPN 100 epoch as example.

bash ./tools/SoCo_FPN_100ep.sh

Finetune detector

1. Copy the folder detectron2_configs to the root folder of Detectron2
2. Train the detectors with Detectron2

Citation

@article{wei2021aligning,
  title={Aligning Pretraining for Detection via Object-Level Contrastive Learning},
  author={Wei, Fangyun and Gao, Yue and Wu, Zhirong and Hu, Han and Lin, Stephen},
  journal={arXiv preprint arXiv:2106.02637},
  year={2021}
}