moco-v3:基于 PyTorch 的自监督 ResNet 与 ViT 实现项目

自监督ResNet与ViT的MoCo v3实现

简介

这是一个用于自监督ResNet和ViT的PyTorch版MoCo v3实现。

原始的MoCo v3是用Tensorflow编写的，并在TPU上运行。此仓库重新实现了该方法并适配了PyTorch和GPU环境。尽管存在库和数值上的差异，但本仓库重现了论文中的结果和观察。

主要成果

以下结果基于ImageNet-1k的自监督预训练，随后是对线性评估或端到端微调的ImageNet-1k监督训练。所有表中所示的结果都是基于批量大小为4096的。

预训练模型和配置可以在CONFIG.md找到。

ResNet-50，线性分类

预训练 轮数	预训练 裁剪方式	线性 准确率
100	2x224	68.9
300	2x224	72.8
1000	2x224	74.6

ViT，线性分类

模型	预训练 轮数	预训练 裁剪方式	线性 准确率
ViT-Small	300	2x224	73.2
ViT-Base	300	2x224	76.7

ViT，端到端微调

模型	预训练 轮数	预训练 裁剪方式	端到端 准确率
ViT-Small	300	2x224	81.4
ViT-Base	300	2x224	83.2

端到端微调的结果使用DeiT仓库得到，采用DeiT的所有默认配置。ViT-B微调150轮（DeiT-B的300轮有81.8%的准确性）。

使用：准备

安装PyTorch，并按照官方PyTorch ImageNet训练代码下载ImageNet数据集。类似于MoCo v1/2，此仓库对官方的PyTorch ImageNet代码进行了最小修改。我们假设用户能够成功运行官方的PyTorch ImageNet代码。 对于ViT模型，需要安装timm（版本为0.4.9）。

已测试CUDA 10.2/CuDNN 7.6.5、PyTorch 1.9.0和timm 0.4.9环境下代码正常运行。

使用：自监督预训练

以下是三个MoCo v3预训练的例子。

两节点(16-GPU)训练ResNet-50，批量大小4096

在第一个节点上运行：

python main_moco.py \
  --moco-m-cos --crop-min=.2 \
  --dist-url 'tcp://[你的第一个节点地址]:[指定端口]' \
  --multiprocessing-distributed --world-size 2 --rank 0 \
  [包含train和val文件夹的ImageNet目录]

在第二个节点上运行相同命令，但--rank改为1。 批量大小为4096时，训练可以适应2个节点共计16个Volta 32G GPU。

单节点(8-GPU)训练ViT-Small，批量大小1024

python main_moco.py \
  -a vit_small -b 1024 \
  --optimizer=adamw --lr=1.5e-4 --weight-decay=.1 \
  --epochs=300 --warmup-epochs=40 \
  --stop-grad-conv1 --moco-m-cos --moco-t=.2 \
  --dist-url 'tcp://localhost:10001' \
  --multiprocessing-distributed --world-size 1 --rank 0 \
  [包含train和val文件夹的ImageNet目录]

八节点训练ViT-Base，批量大小4096

以批量大小4096训练ViT-Base时，使用8个节点：

python main_moco.py \
  -a vit_base \
  --optimizer=adamw --lr=1.5e-4 --weight-decay=.1 \
  --epochs=300 --warmup-epochs=40 \
  --stop-grad-conv1 --moco-m-cos --moco-t=.2 \
  --dist-url 'tcp://[你的第一个节点地址]:[指定端口]' \
  --multiprocessing-distributed --world-size 8 --rank 0 \
  [包含train和val文件夹的ImageNet目录]

其他节点上，分别运行相同的命令，但--rank改为1至7。

注意事项：

1. -b指定的批量大小是所有GPU总和。
2. --lr指定的是基本学习率，会在这一行根据线性学习率缩放规则调整。
3. 较小的批量大小会有更稳定的结果（见论文），但速度较慢。在TPU上（如我们在论文中所做）使用大批量尺寸对于获得良好的速度至关重要。
4. 本仓库仅支持多GPU和DistributedDataParallel训练；不支持单GPU或DataParallel训练。这个代码改进了以更好地适应多节点设置，默认使用自动混合精度进行预训练。

使用方法：端到端微调ViT

为了对ViT进行端到端的微调，使用我们的脚本将其预训练检查点转换为DEiT格式：

python convert_to_deit.py \
  --input [你的检查点路径]/[你的检查点文件].pth.tar \
  --output [目标检查点文件].pth

然后，在DEiT仓库中使用转换后的检查点运行训练：

python $DEIT_DIR/main.py \
  --resume [目标检查点文件].pth \
  --epochs 150

这样，通过150轮微调，ViT-Base模型可以达到83.2%的准确率。

注意：

1. 在DEiT仓库中，我们使用--resume而非--finetune，因为其--finetune选项会在评估模式下训练。加载预训练模型时，需修改为model_without_ddp.load_state_dict(checkpoint['model'], strict=False)。
2. 我们的ViT-Small在Transformer块中使用了heads=12，而在DEiT默认配置中是heads=6。如果要微调我们的ViT-Small模型，请相应地修改DEiT代码。

模型配置

请参考CONFIG.md文档中的命令，了解特定模型配置以及推荐的超参数和预训练基准模型。

迁移学习

请参阅transfer目录中的指南。

许可证

该项目遵循CC-BY-NC 4.0许可证。有关详情，请查看LICENSE文件。

引用

@Article{chen2021mocov3,
  author  = {Xinlei Chen* and Saining Xie* and Kaiming He},
  title   = {An Empirical Study of Training Self-Supervised Vision Transformers},
  journal = {arXiv preprint arXiv:2104.02057},
  year    = {2021},
}