掩码自编码器是可扩展的视觉学习器 的非官方 PyTorch 实现

本仓库基于 BEiT 构建，非常感谢！

目前，我们已按照论文实现了预训练和微调流程，但仍然无法保证能够复现论文中报告的性能！

差异

shuffle 和 unshuffle

在 PyTorch 中似乎无法直接使用 shuffle 和 unshuffle 操作，因此我们采用了另一种方法来实现这一过程：

• 对于 shuffle，我们沿用了 BEiT 中随机生成 14x14 掩码图的方法，其中 mask=0 表示保留该 token，mask=1 表示丢弃该 token（不参与编码器计算）。然后将所有可见 token（mask=0）输入到编码器网络中。
• 对于 unshuffle，我们根据掩码图获取所有被掩码 token 的位置嵌入（并加入共享的掩码 token），再将其与来自编码器的可见 token 拼接在一起，作为解码器的输入进行重建。

正弦-余弦位置嵌入

论文中提到的位置嵌入是 正弦-余弦 版本。我们采用了 此处 的实现，但看起来像是 1D 的嵌入，而非 2D 的。因此，尚不清楚这会带来怎样的影响。另外，我在 MoCoV3 中找到了 2D 的正弦-余弦位置嵌入。如果有兴趣的话，可以尝试一下。

待办事项

• 实现微调流程
• 复用 modeling_pretrain.py 中的模型
• 计算归一化像素目标
• 在编码器中添加 cls token
• 可视化重建图像
• KNN 和线性探测
• …

环境配置

pip install -r requirements.txt

运行

1. 预训练

# 设置保存检查点的路径
OUTPUT_DIR='output/pretrain_mae_base_patch16_224'
# ImageNet-1K 训练集路径
DATA_PATH='/path/to/ImageNet_ILSVRC2012/train'


# 根据显卡情况调整 batch_size
OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node=8 run_mae_pretraining.py \
        --data_path ${DATA_PATH} \
        --mask_ratio 0.75 \
        --model pretrain_mae_base_patch16_224 \
        --batch_size 128 \
        --opt adamw \
        --opt_betas 0.9 0.95 \
        --warmup_epochs 40 \
        --epochs 1600 \
        --output_dir ${OUTPUT_DIR}

2. 微调

# 设置保存检查点的路径
OUTPUT_DIR='output/'
# ImageNet-1K 数据集路径
DATA_PATH='/path/to/ImageNet_ILSVRC2012'
# 预训练模型路径
MODEL_PATH='/path/to/pretrain/checkpoint.pth'

# 根据显卡情况调整 batch_size
OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node=8 run_class_finetuning.py \
    --model vit_base_patch16_224 \
    --data_path ${DATA_PATH} \
    --finetune ${MODEL_PATH} \
    --output_dir ${OUTPUT_DIR} \
    --batch_size 128 \
    --opt adamw \
    --opt_betas 0.9 0.999 \
    --weight_decay 0.05 \
    --epochs 100 \
    --dist_eval

3. 可视化重建

# 设置保存图像的路径
OUTPUT_DIR='output/'
# 用于可视化图像的路径
IMAGE_PATH='files/ILSVRC2012_val_00031649.JPEG'
# 预训练模型路径
MODEL_PATH='/path/to/pretrain/checkpoint.pth'

# 目前仅支持具有归一化像素目标的预训练模型
python run_mae_vis.py ${IMAGE_PATH} ${OUTPUT_DIR} ${MODEL_PATH}

结果

模型	预训练轮数	微调轮数	准确率	日志	权重链接
vit-base	400轮	100轮	83.1%	预训练 微调	Google Drive 百度网盘(提取码: mae6)
vit-large	400轮	50轮	84.5%	预训练 微调	未提供

由于 GPU 资源有限，要按照论文中的更大模型或更长调度进行预训练确实颇具挑战。（vit-large 模型的预训练及端到端微调是由一位热情的帅哥完成的，他使用了多张 V100 显卡，但权重并未公开）

因此，如果有人能够成功完成，请随时在 issue 中报告或提交 PR，非常感谢！

你的 Star 就是我的动力，谢谢！

MAE-pytorch 快速上手指南

MAE-pytorch 是论文《Masked Autoencoders Are Scalable Vision Learners》的非官方 PyTorch 实现，基于 BEiT 代码库构建。本项目实现了预训练（Pretrain）和微调（Finetune）流程，并支持重建图像可视化。

注意：由于硬件限制及实现细节差异（如位置编码），本仓库无法保证完全复现原论文报告的性能指标。

1. 环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: 3.6+
• PyTorch: 兼容版本（建议 1.8+）
• GPU: 支持 CUDA 的 NVIDIA 显卡（多卡训练效果更佳）

前置依赖

确保已安装基础深度学习环境。项目所需的具体依赖包如下：

• torch, torchvision
• timm
• Pillow, numpy, matplotlib 等常规科学计算库

2. 安装步骤

克隆仓库并安装依赖：

git clone https://github.com/pengzhiliang/MAE-pytorch.git
cd MAE-pytorch
pip install -r requirements.txt

国内加速建议：如果下载依赖较慢，可使用清华或阿里镜像源：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 基本使用

使用前请准备好 ImageNet-1k 数据集路径。以下示例展示最核心的三个步骤：预训练、微调和可视化。

3.1 预训练 (Pretrain)

设置检查点保存路径和数据集路径，启动分布式预训练。

# 设置变量
OUTPUT_DIR='output/pretrain_mae_base_patch16_224'
DATA_PATH='/path/to/ImageNet_ILSVRC2012/train'

# 启动预训练 (示例使用 8 张卡，batch_size 可根据显存调整)
OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node=8 run_mae_pretraining.py \
        --data_path ${DATA_PATH} \
        --mask_ratio 0.75 \
        --model pretrain_mae_base_patch16_224 \
        --batch_size 128 \
        --opt adamw \
        --opt_betas 0.9 0.95 \
        --warmup_epochs 40 \
        --epochs 1600 \
        --output_dir ${OUTPUT_DIR}

3.2 微调 (Finetune)

加载预训练权重进行分类任务微调。

# 设置变量
OUTPUT_DIR='output/'
DATA_PATH='/path/to/ImageNet_ILSVRC2012'
MODEL_PATH='/path/to/pretrain/checkpoint.pth'

# 启动微调
OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node=8 run_class_finetuning.py \
    --model vit_base_patch16_224 \
    --data_path ${DATA_PATH} \
    --finetune ${MODEL_PATH} \
    --output_dir ${OUTPUT_DIR} \
    --batch_size 128 \
    --opt adamw \
    --opt_betas 0.9 0.999 \
    --weight_decay 0.05 \
    --epochs 100 \
    --dist_eval

3.3 重建可视化 (Visualization)

查看模型对掩码图像的重建效果（仅支持使用归一化像素目标训练的模型）。

# 设置变量
OUTPUT_DIR='output/'
IMAGE_PATH='files/ILSVRC2012_val_00031649.JPEG'
MODEL_PATH='/path/to/pretrain/checkpoint.pth'

# 运行可视化脚本
python run_mae_vis.py ${IMAGE_PATH} ${OUTPUT_DIR} ${MODEL_PATH}

预训练模型参考

若需直接使用预训练权重，可参考以下结果（ViT-Base, 400 epoch 预训练 + 100 epoch 微调，准确率约 83.1%）：

• Google Drive: 下载链接
• 百度网盘: 下载链接 (提取码: mae6)