密集对比学习用于自监督视觉预训练

本项目托管了实现DenseCL算法进行自监督表示学习的代码。

密集对比学习用于自监督视觉预训练,
王新龙, 张汝峰, 沈春华, 孔涛, 李磊
发表于: 2021年IEEE计算机视觉与模式识别会议(CVPR), 口头报告
arXiv预印本 (arXiv 2011.09157)

亮点2

亮点

提升密集预测效果: DenseCL预训练模型极大地提升了密集预测任务的性能,包括目标检测和语义分割(分别提高至多2%的AP和3%的mIoU)。
简单实现: DenseCL的核心部分可以用10行代码实现,因此易于使用和修改。
灵活使用: DenseCL与数据预处理解耦,因此能够快速灵活地训练,同时不受使用何种数据增强方式和图像采样方式的影响。
高效训练: 与基线方法相比,我们的方法引入的计算开销可以忽略不计(仅慢<1%)。

更新

提供了简单教程来在AdelaiDet中使用DenseCL(例如,与SOLOv2和FCOS一起使用)。(2021/05/16)
发布了DenseCL的代码和预训练模型。(2021/02/03)

安装

请参考INSTALL.md进行安装和数据集准备。

模型

为了方便使用,我们提供了以下在COCO或ImageNet上预训练的模型。

预训练方法	预训练数据集	骨干网络	训练轮数	训练时间	VOC检测	VOC分割	链接
MoCo-v2	COCO	ResNet-50	800	1.0天	54.7	64.5
DenseCL	COCO	ResNet-50	800	1.0天	56.7	67.5	下载
DenseCL	COCO	ResNet-50	1600	2.0天	57.2	68.0	下载
MoCo-v2	ImageNet	ResNet-50	200	2.3天	57.0	67.5
DenseCL	ImageNet	ResNet-50	200	2.3天	58.7	69.4	下载
DenseCL	ImageNet	ResNet-101	200	4.3天	61.3	74.1	下载

注意:

VOC检测和分割的指标分别为AP(COCO风格)和mIoU。结果是5次试验的平均值。
训练时间是在8个V100 GPU上测量的。
有关不同基准测试的更多结果,请参阅我们的论文。

我们还提供了在AdelaiDet模型中使用DenseCL的实验,例如SOLOv2和FCOS。请参考使用说明获取简单用法。

SOLOv2在COCO实例分割上的结果

预训练方法	预训练数据集	mask AP
有监督	ImageNet	35.2
MoCo-v2	ImageNet	35.2
DenseCL	ImageNet	35.7 (+0.5)

FCOS在COCO目标检测上的结果

预训练方法	预训练数据集	box AP
有监督	ImageNet	39.9
MoCo-v2	ImageNet	40.3
DenseCL	ImageNet	40.9 (+1.0)

使用方法

训练

./tools/dist_train.sh configs/selfsup/densecl/densecl_coco_800ep.py 8

提取骨干网络权重

WORK_DIR=work_dirs/selfsup/densecl/densecl_coco_800ep/
CHECKPOINT=${WORK_DIR}/epoch_800.pth
WEIGHT_FILE=${WORK_DIR}/extracted_densecl_coco_800ep.pth

python tools/extract_backbone_weights.py ${CHECKPOINT} ${WEIGHT_FILE}

迁移到目标检测和分割

请参考README.md进行目标检测和语义分割的迁移。请参考使用说明将其迁移到AdelaiDet中的密集预测模型,例如SOLOv2和FCOS。

提示

提取骨干网络权重后,该模型可用于替代原始ImageNet预训练模型,作为许多密集预测任务的初始化。
如果您的机器存在数据加载速度慢的问题,特别是对于ImageNet,建议您通过folder2lmdb_imagenet.py或folder2lmdb_coco.py将ImageNet转换为lmdb格式,并使用此config_imagenet或config_coco进行训练。

致谢

我们要感谢OpenSelfSup的开源项目和PyContrast提供的检测评估配置。

引用

如果本项目对您的研究有帮助,请考虑在您的出版物中引用我们的论文。BibTeX参考如下。

@inproceedings{wang2020DenseCL,
  title={Dense Contrastive Learning for Self-Supervised Visual Pre-Training},
  author={Wang, Xinlong and Zhang, Rufeng and Shen, Chunhua and Kong, Tao and Li, Lei},
  booktitle =  {Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR)},
  year={2021}
}