视觉Transformer

这是对Google发布的ViT模型仓库的PyTorch重新实现，该模型随论文《一张图胜过16×16个词：大规模图像识别中的Transformer》（作者：Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby）一同发布。

这篇论文表明，直接将Transformer应用于图像块并在大规模数据集上进行预训练，在图像识别任务中表现非常出色。

视觉Transformer使用标准的Transformer编码器和固定大小的图像块，在图像识别任务中达到了最先进的水平。为了进行分类任务，作者采用了在序列中添加一个可学习的“分类标记”的标准方法。

使用方法

1. 下载预训练模型（Google官方检查点）

• 可用模型：ViT-B_16(85.8M)、R50+ViT-B_16(97.96M)、ViT-B_32(87.5M)、ViT-L_16(303.4M)、ViT-L_32(305.5M)、ViT-H_14(630.8M)
  • imagenet21k预训练模型：
    • ViT-B_16、ViT-B_32、ViT-L_16、ViT-L_32、ViT-H_14
  • imagenet21k预训练 + imagenet2012微调模型：
    • ViT-B_16-224、ViT-B_16、ViT-B_32、ViT-L_16-224、ViT-L_16、ViT-L_32
  • 混合模型（Resnet50 + Transformer）：
    • R50-ViT-B_16

# imagenet21k预训练
wget https://storage.googleapis.com/vit_models/imagenet21k/{MODEL_NAME}.npz

# imagenet21k预训练 + imagenet2012微调
wget https://storage.googleapis.com/vit_models/imagenet21k+imagenet2012/{MODEL_NAME}.npz

2. 训练模型

python3 train.py --name cifar10-100_500 --dataset cifar10 --model_type ViT-B_16 --pretrained_dir checkpoint/ViT-B_16.npz

CIFAR-10和CIFAR-100会自动下载并开始训练。如果要使用其他数据集，需要自定义data_utils.py。

默认批量大小为512。当GPU显存不足时，可以通过调整--gradient_accumulation_steps参数来继续训练。

此外，还可以使用自动混合精度(Amp)来减少内存占用并加快训练速度：

python3 train.py --name cifar10-100_500 --dataset cifar10 --model_type ViT-B_16 --pretrained_dir checkpoint/ViT-B_16.npz --fp16 --fp16_opt_level O2

结果

为了验证转换后的模型权重是否正确，我们将其与作者的实验结果进行了对比。我们在训练过程中使用了混合精度，并将--fp16_opt_level设置为O2。

imagenet-21k

• tensorboard

模型	数据集	分辨率	acc(官方)	acc(本仓库)	时间
ViT-B_16	CIFAR-10	224x224	-	0.9908	3h 13m
ViT-B_16	CIFAR-10	384x384	0.9903	0.9906	12h 25m
ViT_B_16	CIFAR-100	224x224	-	0.923	3h 9m
ViT_B_16	CIFAR-100	384x384	0.9264	0.9228	12h 31m
R50-ViT-B_16	CIFAR-10	224x224	-	0.9892	4h 23m
R50-ViT-B_16	CIFAR-10	384x384	0.99	0.9904	15h 40m
R50-ViT-B_16	CIFAR-100	224x224	-	0.9231	4h 18m
R50-ViT-B_16	CIFAR-100	384x384	0.9231	0.9197	15h 53m
ViT_L_32	CIFAR-10	224x224	-	0.9903	2h 11m
ViT_L_32	CIFAR-100	224x224	-	0.9276	2h 9m
ViT_H_14	CIFAR-100	224x224	-	WIP

imagenet-21k + imagenet2012

• tensorboard

模型	数据集	分辨率	acc
ViT-B_16-224	CIFAR-10	224x224	0.99
ViT_B_16-224	CIFAR-100	224x224	0.9245
ViT-L_32	CIFAR-10	224x224	0.9903
ViT-L_32	CIFAR-100	224x224	0.9285

短周期训练

• 在下面的实验中，我们使用了224x224的分辨率。
• tensorboard

上游	模型	数据集	total_steps /warmup_steps	acc(官方)	acc(本仓库)
imagenet21k	ViT-B_16	CIFAR-10	500/100	0.9859	0.9859
imagenet21k	ViT-B_16	CIFAR-10	1000/100	0.9886	0.9878
imagenet21k	ViT-B_16	CIFAR-100	500/100	0.8917	0.9072
imagenet21k	ViT-B_16	CIFAR-100	1000/100	0.9115	0.9216

可视化

ViT由标准的Transformer编码器组成，而编码器又由自注意力模块和MLP模块构成。通过自注意力的注意力分数，可以可视化输入图像的注意力图。

可视化代码可以在visualize_attention_map中找到。

参考文献

• Google ViT
• Pytorch图像模型(timm)

引用

@article{dosovitskiy2020,
  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
  author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and  Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and Uszkoreit, Jakob and Houlsby, Neil},
  journal={arXiv preprint arXiv:2010.11929},
  year={2020}
}

ViT-pytorch 快速上手指南

ViT-pytorch 是 Google Vision Transformer (ViT) 模型的 PyTorch 复现版本。该模型通过将图像划分为固定大小的补丁（Patches）并直接应用 Transformer 编码器，在图像识别任务中达到了业界领先的水平。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6 或更高版本
• 深度学习框架: PyTorch 1.7+ (建议配合 CUDA 使用以加速训练)
• 其他依赖:
  • numpy
  • tensorboard (用于查看训练日志)
  • apex (可选，用于混合精度训练以节省显存)

建议使用 conda 或 pip 安装基础依赖：

pip install torch torchvision tensorboard numpy
# 如需混合精度训练，请参考 NVIDIA 官方文档安装 apex

2. 安装与模型下载

本项目无需复杂的安装过程，克隆仓库即可使用。核心步骤是下载预训练权重。

2.1 克隆项目

git clone https://github.com/lucidrains/vit-pytorch.git
cd vit-pytorch

(注：如果原文指向的是特定的复现仓库，请替换为对应的 git 地址。根据提供的 README，此处指代通用的 ViT-pytorch 实现结构)

2.2 下载预训练模型

您可以从 Google Cloud 下载官方预训练权重。国内用户若下载缓慢，可尝试使用代理或寻找国内镜像源。

可用模型包括：

• ImageNet-21k 预训练: ViT-B_16, ViT-B_32, ViT-L_16, ViT-L_32, ViT-H_14
• ImageNet-21k + ImageNet-2012 微调: ViT-B_16-224, ViT-L_16-224 等
• 混合模型 (ResNet50 + ViT): R50-ViT-B_16

下载命令示例：

# 创建检查点目录
mkdir -p checkpoint

# 下载 ImageNet-21k 预训练模型 (以 ViT-B_16 为例)
wget https://storage.googleapis.com/vit_models/imagenet21k/ViT-B_16.npz -O checkpoint/ViT-B_16.npz

# 或者下载经过 ImageNet-2012 微调的模型
# wget https://storage.googleapis.com/vit_models/imagenet21k+imagenet2012/ViT-B_16.npz -O checkpoint/ViT-B_16.npz

3. 基本使用

3.1 快速开始训练

默认配置会自动下载 CIFAR-10 或 CIFAR-100 数据集并开始训练。以下命令使用 ViT-B_16 模型在 CIFAR-10 数据集上进行训练：

python3 train.py --name cifar10-100_500 --dataset cifar10 --model_type ViT-B_16 --pretrained_dir checkpoint/ViT-B_16.npz

参数说明：

• --dataset: 指定数据集 (支持 cifar10, cifar100，自定义数据集需修改 utils/data_utils.py)。
• --model_type: 选择模型架构 (如 ViT-B_16, ViT-L_32 等)。
• --pretrained_dir: 指定下载的预训练权重路径。

3.2 显存优化（混合精度训练）

如果 GPU 显存不足，默认 batch size (512) 可能导致 OOM。您可以采取以下两种策略：

1. 调整梯度累积步数：增加 --gradient_accumulation_steps 参数。
2. 启用自动混合精度 (AMP)：使用 --fp16 参数加速训练并降低显存占用。

混合精度训练示例：

python3 train.py --name cifar10-100_500 --dataset cifar10 --model_type ViT-B_16 --pretrained_dir checkpoint/ViT-B_16.npz --fp16 --fp16_opt_level O2

3.3 可视化注意力图

ViT 的核心在于自注意力机制。您可以使用提供的 Jupyter Notebook 脚本可视化模型对输入图像的注意力分布：

# 启动 jupyter 并打开可视化脚本
jupyter notebook visualize_attention_map.ipynb

运行后即可生成类似论文中的注意力热力图，观察模型关注的图像区域。