由于可用资源有限，我们仅在 cifar10 上测试模型。我们主要想重现这样的结果：使用 MAE 预训练 ViT 可以比直接使用标签进行监督学习训练获得更好的结果。这应该是自我监督学习比监督学习更有效的数据的证据。
我们主要遵循论文中的实现细节。但是，由于 Cifar10 和 ImageNet 的区别，我们做了一些修改：
- 我们使用 vit-tiny 而不是 vit-base。
- 由于 Cifar10 只有 50k 训练数据，我们将 pretraining epoch 从 400 增加到 2000，将 warmup epoch 从 40 增加到 200。我们注意到，在 2000 epoch 之后损失仍在减少。
- 我们将训练分类器的批量大小从 1024 减少到 512 以减轻过度拟合。
pip install -r requirements.txt
首先进行预训练
# pretrained with mae
python mae_pretrain.py训练未用MAE的分类器，也就是从头开始训练分类器
# train classifier from scratch
python train_classifier.py
利用训练好的MAE的encoder作为输入，构建的分类模型作为分类器
# train classifier from pretrained model
python train_classifier.py --pretrained_model_path vit-t-mae.pth --output_model_path vit-t-classifier-from_pretrained.pth集成了tensorboerd
tensorboard --logdir logs
可以查看结果
| Model | Validation Acc |
|---|---|
| ViT-T w/o pretrain | 74.13 |
| ViT-T w/ pretrain | 89.77 |
可视化CIFAR10前16张的图片，也可以在TensorBoard中查看
由于TensorBoard的文件太大了，所以没有传上去，不过我上传了，大家可以自由查看https://tensorboard.dev/experiment/GIv9UzukQ5yFalbykPVMFQ/
其他的所有的训练权重和logs，最后都上传到release上了，可以自取
This project is under the MIT license. See LICENSE for details.
该项目主要基于 https://github.com/IcarusWizard/MAE 进行了一些翻译和注释，同时还加入了其他的一些功能。