密集对比学习用于自监督视觉预训练
本项目托管了实现DenseCL算法进行自监督表示学习的代码。
密集对比学习用于自监督视觉预训练,
王新龙, 张汝峰, 沈春华, 孔涛, 李磊
发表于: 2021年IEEE计算机视觉与模式识别会议(CVPR), 口头报告
arXiv预印本 (arXiv 2011.09157)
亮点
- 提升密集预测效果: 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}
}