MAE 在 HuggingFace Transformers 中可用(支持 PyTorch 和 TF)
大家好!感谢你们的出色工作。几个月以来,MAE 在 HuggingFace Transformers 中提供了易于使用且熟悉的 API: https://huggingface.co/docs/transformers/model_doc/vit_mae 该模型既可在 PyTorch 中使用,也可在 TensorFlow 中使用。API 与 BertModel 相同(但称为 ViTMAEModel)。还有一个 ViTMAEForPreTraining 类,与 BertForPreTraining 类类似,其中包括解码器和损失计算。我们提供了一个易于使用的脚本,可用于直接在自定义数据上预训练 ViTMAE 模型,该脚本名为 run_mae.py。可以在以下位置找到该脚本: https://GitHub.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining。该脚本利用了 Trainer 并直接在分布式环境中运行。该脚本还支持在训练期间和/或之后将模型推送到 hub(使用 from_pretrained 后可轻松重新加载)。预训练后,您可以将权重自动加载到 ViTForImageClassification 模型中,从而实现微调。微调 notebook 可在 此处 (针对 Torchvision 的数据增强)或 此处 (如果您更喜欢使用 Albumentations 进行数据增强)找到。我还创建了一个 notebook,用于演示可视化(基于此仓库): https://GitHub.com/NielsRogge/Transformers-Tutorials/blob/master/ViTMAE/ViT_MAE_visualization_demo.ipynb
内容来源: facebookresearch/mae