视觉Transformer与MLP-Mixer架构

在本仓库中，我们发布了来自以下论文的模型：

• 《一张图像胜过16×16个词：大规模图像识别中的Transformer》（https://arxiv.org/abs/2010.11929）
• 《MLP-Mixer：一种全MLP的视觉架构》（https://arxiv.org/abs/2105.01601）
• 《如何训练你的ViT？视觉Transformer中的数据、增强与正则化》（https://arxiv.org/abs/2106.10270）
• 《当视觉Transformer无需预训练或强数据增强即可超越ResNet时》（https://arxiv.org/abs/2106.01548）
• 《LiT：基于锁定图像文本微调的零样本迁移》（https://arxiv.org/abs/2111.07991）
• 《代理差距最小化改进了尖锐度感知训练》（https://arxiv.org/abs/2203.08065）

这些模型已在ImageNet和ImageNet-21k数据集上进行了预训练。我们提供了使用JAX/Flax对已发布模型进行微调的代码。

本代码库中的模型最初是在https://github.com/google-research/big_vision/中训练的，在那里您可以找到更高级的代码（例如多主机训练），以及一些原始的训练脚本（例如用于预训练ViT的[configs/vit_i21k.py](https://github.com/google-research/big_vision/blob/main/big_vision/configs/vit_i21k.py)，或用于模型迁移的[configs/transfer.py](https://github.com/google-research/big_vision/blob/main/big_vision/configs/transfer.py)）。

目录：

• 视觉Transformer与MLP-Mixer架构
  • Colab
  • 安装
  • 微调模型
  • 视觉Transformer
    • 可用的ViT模型
    • 预期的ViT结果
  • MLP-Mixer
    • 可用的Mixer模型
    • 预期的Mixer结果
  • LiT模型
  • 云端运行
    • 创建虚拟机
    • 设置虚拟机
  • BibTeX
  • 免责声明
  • 更新日志

Colab

以下Colab既可以使用GPU，也可以使用TPU（8核，数据并行）运行。

第一个Colab展示了视觉Transformer和MLP Mixer的JAX代码。该Colab允许您直接在Colab界面中编辑仓库中的文件，并配有注释过的单元格，逐步引导您理解代码，同时让您与数据进行交互。

https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax.ipynb

第二个Colab可以让您探索超过5万个用于生成第三篇论文“如何训练你的ViT？…”数据的视觉Transformer及混合模型检查点。该Colab包含用于探索和选择检查点的代码，以及使用本仓库的JAX代码和流行的timm PyTorch库进行推理的功能——后者可以直接加载这些检查点。需要注意的是，少数模型也可直接从TF-Hub获取：sayakpaul/collections/vision_transformer（由Sayak Paul贡献）。

第二个Colab还允许您在任何tfds数据集以及您自己的数据集上对检查点进行微调，这些数据集可以是单独的JPEG文件形式（可选地直接从Google Drive读取）。

https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax_augreg.ipynb

注意：截至2021年6月20日，Google Colab仅支持单个GPU（Nvidia Tesla T4），而TPU（目前为TPUv2-8）则是间接连接到Colab虚拟机并通过较慢的网络通信，这导致训练速度非常慢。如果您有大量数据需要微调，通常建议搭建一台专用机器。有关详细信息，请参阅云端运行部分。

安装

请确保您的机器上已安装Python>=3.10。

通过运行以下命令安装JAX和Python依赖项：

# 如果使用GPU：
pip install -r vit_jax/requirements.txt

# 如果使用TPU：
pip install -r vit_jax/requirements-tpu.txt

对于较新版本的JAX，请遵循此处链接的相应仓库提供的说明。请注意，CPU、GPU和TPU的安装说明略有不同。

安装Flaxformer，并按照此处链接的相应仓库提供的说明进行操作。

更多详情请参阅下方的云端运行部分。

微调模型

您可以在感兴趣的自己的数据集上对下载的模型进行微调。所有模型都使用相同的命令行界面。

例如，要在 CIFAR10 数据集上微调一个在 ImageNet21k 上预训练的 ViT-B/16 模型（注意我们如何将 b16,cifar10 作为配置参数传入，以及如何指示代码直接从 GCS 存储桶中加载模型，而不是先将其下载到本地目录）：

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 \
    --config.pretrained_dir='gs://vit_models/imagenet21k'

如果要对在 ImageNet21k 上预训练的 Mixer-B/16 模型进行 CIFAR10 数据集上的微调：

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/mixer_base16_cifar10.py \
    --config.pretrained_dir='gs://mixer_models/imagenet21k'

论文《如何训练你的 ViT？》添加了超过 5 万个检查点，您可以使用 configs/augreg.py 配置对其进行微调。当您仅指定模型名称（即 configs/model.py 中的 config.name 值）时，系统会自动选择上游验证准确率最高的 i21k 检查点（“推荐”检查点，详见论文第 4.5 节）。要决定使用哪个模型，可以参考论文中的图 3。此外，您也可以选择其他检查点（参见 Colab 笔记本 vit_jax_augreg.ipynb），并指定来自 filename 或 adapt_filename 列的值，这些值对应于 gs://vit_models/augreg 目录中不带 .npz 后缀的文件名。

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/augreg.py:R_Ti_16 \
    --config.dataset=oxford_iiit_pet \
    --config.base_lr=0.01

目前，代码会自动下载 CIFAR-10 和 CIFAR-100 数据集。其他公开或自定义数据集也可以通过 TensorFlow Datasets 库 轻松集成。请注意，您还需要更新 vit_jax/input_pipeline.py 文件，以指定新增数据集的相关参数。

需要注意的是，我们的代码会利用所有可用的 GPU/TPU 进行微调。

如需查看所有可用标志的详细列表，请运行 python3 -m vit_jax.train --help。

内存注意事项：

• 不同模型所需的内存不同。可用内存还取决于加速器的配置（类型和数量）。如果遇到内存不足错误，您可以将 --config.accum_steps=8 的值调高；或者相应地降低 --config.batch=512（并按比例调整 --config.base_lr）。
• 主机端会在内存中维护一个打乱缓冲区。如果您遇到主机端 OOM 错误（而非加速器端 OOM），可以将默认的 --config.shuffle_buffer=50000 值调低。

视觉 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 编码器。为了进行分类，我们采用在序列中添加一个额外的可学习“分类标记”的标准方法。

可用的 ViT 模型

我们在不同的 GCS 存储桶中提供了多种 ViT 模型。例如，可以通过以下命令下载模型：

wget https://storage.googleapis.com/vit_models/imagenet21k/ViT-B_16.npz

模型文件名（不含 .npz 后缀）与 vit_jax/configs/models.py 中的 config.model_name 对应。

• gs://vit_models/imagenet21k：在 ImageNet-21k 上预训练的模型。
• gs://vit_models/imagenet21k+imagenet2012：在 ImageNet-21k 上预训练并在 ImageNet 上微调的模型。
• gs://vit_models/augreg：在 ImageNet-21k 上预训练并应用不同程度 AugReg 技术的模型，性能有所提升。
• gs://vit_models/sam：使用 SAM 技术在 ImageNet 上预训练的模型。
• gs://vit_models/gsam：使用 GSAM 技术在 ImageNet 上预训练的模型。

我们推荐使用以下经过 AugReg 训练且预训练指标最优的检查点：

模型	预训练检查点	大小	微调后检查点	分辨率	图像/秒	ImageNet 准确率
L/16	gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0.npz	1243 MiB	gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz	384	50	85.59%
B/16	gs://vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz	391 MiB	gs://vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz	384	138	85.49%
S/16	gs://vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz	115 MiB	gs://vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz	384	300	83.73%
R50+L/32	gs://vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1.npz	1337 MiB	gs://vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz	384	327	85.99%
R26+S/32	gs://vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0.npz	170 MiB	gs://vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz	384	560	83.85%
Ti/16	gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz	37 MiB	gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz	384	610	78.22%
B/32	gs://vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0.npz	398 MiB	gs://vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz	384	955	83.59%
S/32	gs://vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0.npz	118 MiB	gs://vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz	384	2154	79.58%
R+Ti/16	gs://vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz	40 MiB	gs://vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz	384	2426	75.40%

使用 gs://vit_models/imagenet21k 中的模型，已复现了原始 ViT 论文（https://arxiv.org/abs/2010.11929）中的结果：

模型	数据集	dropout=0.0	dropout=0.1
R50+ViT-B_16	cifar10	98.72%, 3.9小时 (A100), tb.dev	98.94%, 10.1小时 (V100), tb.dev
R50+ViT-B_16	cifar100	90.88%, 4.1小时 (A100), tb.dev	92.30%, 10.1小时 (V100), tb.dev
R50+ViT-B_16	imagenet2012	83.72%, 9.9小时 (A100), tb.dev	85.08%, 24.2小时 (V100), tb.dev
ViT-B_16	cifar10	99.02%, 2.2小时 (A100), tb.dev	98.76%, 7.8小时 (V100), tb.dev
ViT-B_16	cifar100	92.06%, 2.2小时 (A100), tb.dev	91.92%, 7.8小时 (V100), tb.dev
ViT-B_16	imagenet2012	84.53%, 6.5小时 (A100), tb.dev	84.12%, 19.3小时 (V100), tb.dev
ViT-B_32	cifar10	98.88%, 0.8小时 (A100), tb.dev	98.75%, 1.8小时 (V100), tb.dev
ViT-B_32	cifar100	92.31%, 0.8小时 (A100), tb.dev	92.05%, 1.8小时 (V100), tb.dev
ViT-B_32	imagenet2012	81.66%, 3.3小时 (A100), tb.dev	81.31%, 4.9小时 (V100), tb.dev
ViT-L_16	cifar10	99.13%, 6.9小时 (A100), tb.dev	99.14%, 24.7小时 (V100), tb.dev
ViT-L_16	cifar100	92.91%, 7.1小时 (A100), tb.dev	93.22%, 24.4小时 (V100), tb.dev
ViT-L_16	imagenet2012	84.47%, 16.8小时 (A100), tb.dev	85.05%, 59.7小时 (V100), tb.dev
ViT-L_32	cifar10	99.06%, 1.9小时 (A100), tb.dev	99.09%, 6.1小时 (V100), tb.dev
ViT-L_32	cifar100	93.29%, 1.9小时 (A100), tb.dev	93.34%, 6.2小时 (V100), tb.dev
ViT-L_32	imagenet2012	81.89%, 7.5小时 (A100), tb.dev	81.13%, 15.0小时 (V100), tb.dev

我们还想强调，通过较短的训练周期也能获得高质量的结果，并鼓励使用我们代码的用户尝试调整超参数，以在准确性和计算资源之间取得平衡。 下表展示了一些针对 CIFAR-10/100 数据集的例子。

上游	模型	数据集	总步数 / 预热步数	准确率	实际耗时	链接
imagenet21k	ViT-B_16	cifar10	500 / 50	98.59%	17m	tensorboard.dev
imagenet21k	ViT-B_16	cifar10	1000 / 100	98.86%	39m	tensorboard.dev
imagenet21k	ViT-B_16	cifar100	500 / 50	89.17%	17m	tensorboard.dev
imagenet21k	ViT-B_16	cifar100	1000 / 100	91.15%	39m	tensorboard.dev

MLP-Mixer

作者：Ilya Tolstikhin*、Neil Houlsby*、Alexander Kolesnikov*、Lucas Beyer*， Xiaohua Zhai、Thomas Unterthiner、Jessica Yung、Andreas Steiner、Daniel Keysers、 Jakob Uszkoreit、Mario Lucic、Alexey Dosovitskiy。

(*) 共同第一作者。

MLP-Mixer（简称Mixer）由逐块线性嵌入、Mixer层和分类头组成。Mixer层包含一个通道混合的MLP和一个特征混合的MLP，每个MLP由两个全连接层和一个GELU非线性激活函数构成。其他组件包括：残差连接、Dropout以及线性分类头。

安装步骤与上述相同，请参阅安装说明。

可用的Mixer模型

我们提供了在ImageNet和ImageNet-21k数据集上预训练的Mixer-B/16和Mixer-L/16模型。详细信息请参见Mixer论文中的表3。所有模型均可在以下地址找到：

https://console.cloud.google.com/storage/mixer_models/

请注意，这些模型也可直接从TF-Hub获取： sayakpaul/collections/mlp-mixer（由Sayak Paul贡献的外部资源）。

预期的Mixer结果

我们在Google Cloud的四张V100 GPU机器上运行了微调代码，并使用了本仓库中的默认适配参数。以下是结果：

上游	模型	数据集	准确率	实际耗时	链接
ImageNet	Mixer-B/16	cifar10	96.72%	3.0小时	tensorboard.dev
ImageNet	Mixer-L/16	cifar10	96.59%	3.0小时	tensorboard.dev
ImageNet-21k	Mixer-B/16	cifar10	96.82%	9.6小时	tensorboard.dev
ImageNet-21k	Mixer-L/16	cifar10	98.34%	10.0小时	tensorboard.dev

LiT模型

有关详情，请参阅Google AI博客文章 LiT：为图像模型添加语言理解能力， 或阅读CVPR论文“LiT：通过锁定图像文本微调实现零样本迁移” (https://arxiv.org/abs/2111.07991)。

我们发布了一个Transformer B/16-base模型，其在ImageNet上的零样本准确率为72.1%，以及一个L/16-large模型，其在ImageNet上的零样本准确率为75.7%。更多关于这些模型的详细信息，请参阅 LiT模型卡片。

我们提供了一个基于浏览器的演示，其中包含小型文本编码器，可供交互使用（最小的模型甚至可以在现代手机上运行）：

https://google-research.github.io/vision_transformer/lit/

最后，还有一个Colab笔记本，用于使用带有图像和文本编码器的JAX模型：

https://colab.research.google.com/github/google-research/vision_transformer/blob/main/lit.ipynb

需要注意的是，目前上述任何模型都不支持多语言输入，但我们正在开发此类模型，并将在它们可用时更新此仓库。

该仓库仅包含LiT模型的评估代码。训练代码可在big_vision仓库中找到：

https://github.com/google-research/big_vision/tree/main/big_vision/configs/proj/image_text

预期的零样本结果来自model_cards/lit.md（请注意，零样本评估与Colab中的简化评估略有不同）：

模型	B16B_2	L16L
ImageNet 零样本	73.9%	75.7%
ImageNet v2 零样本	65.1%	66.6%
CIFAR100 零样本	79.0%	80.5%
Pets37 零样本	83.3%	83.3%
Resisc45 零样本	25.3%	25.6%
MS-COCO 文字描述到图像检索	51.6%	48.5%
MS-COCO 图像到文字描述检索	31.8%	31.1%

在云端运行

虽然上述Colab对于入门非常有用，但通常您会希望在配备更强大加速器的大型机器上进行训练。

创建虚拟机

您可以通过以下命令在Google Cloud上设置一台带有GPU的虚拟机：

# 设置以下所有命令中使用的变量。
# 注意，项目必须已设置账单。
# 关于拥有GPU的区域列表，请参考
# https://cloud.google.com/compute/docs/gpus/gpu-regions-zones
PROJECT=my-awesome-gcp-project  # 项目必须启用计费。
VM_NAME=vit-jax-vm-gpu
ZONE=europe-west4-b

# 以下配置已在本仓库中测试过。您可以选择其他镜像和机器组合（例如），请参考相应的gcloud命令：
# gcloud compute images list --project ml-images
# gcloud compute machine-types list
# 等等。
gcloud compute instances create $VM_NAME \
    --project=$PROJECT --zone=$ZONE \
    --image=c1-deeplearning-tf-2-5-cu110-v20210527-debian-10 \
    --image-project=ml-images --machine-type=n1-standard-96 \
    --scopes=cloud-platform,storage-full --boot-disk-size=256GB \
    --boot-disk-type=pd-ssd --metadata=install-nvidia-driver=True \
    --maintenance-policy=TERMINATE \
    --accelerator=type=nvidia-tesla-v100,count=8

# 连接到虚拟机（在机器设置并启动后的几分钟内）。
gcloud compute ssh --project $PROJECT --zone $ZONE $VM_NAME

# 使用完毕后停止虚拟机（停止后的虚拟机仅收取存储费用）。
gcloud compute instances stop --project $PROJECT --zone $ZONE $VM_NAME

# 使用完毕后删除虚拟机（这也将删除虚拟机上存储的所有数据）。
gcloud compute instances delete --project $PROJECT --zone $ZONE $VM_NAME

或者，您也可以使用以下类似的命令来设置带有 TPU 挂载的 Cloud 虚拟机（以下命令摘自 TPU 教程）：

PROJECT=my-awesome-gcp-project  # 项目必须启用结算功能。
VM_NAME=vit-jax-vm-tpu
ZONE=europe-west4-a

# 初始设置服务身份时需要执行此操作。
gcloud beta services identity create --service tpu.googleapis.com

# 创建直接挂载 TPU 的虚拟机。
gcloud alpha compute tpus tpu-vm create $VM_NAME \
    --project=$PROJECT --zone=$ZONE \
    --accelerator-type v3-8 \
    --version tpu-vm-base

# 连接到虚拟机（在机器完成设置并启动后的几分钟内）。
gcloud alpha compute tpus tpu-vm ssh --project $PROJECT --zone $ZONE $VM_NAME

# 使用完毕后停止虚拟机（已停止的虚拟机仅按存储收费）。
gcloud alpha compute tpus tpu-vm stop --project $PROJECT --zone $ZONE $VM_NAME

# 使用完毕后删除虚拟机（这也将删除虚拟机上存储的所有数据）。
gcloud alpha compute tpus tpu-vm delete --project $PROJECT --zone $ZONE $VM_NAME

设置虚拟机

然后按照常规方式克隆仓库并安装依赖项（包括支持 TPU 的 jaxlib）：

git clone --depth=1 --branch=master https://github.com/google-research/vision_transformer
cd vision_transformer

# 可选：安装 virtualenv
pip3 install virtualenv
python3 -m virtualenv env
. env/bin/activate

如果您连接的是带有 GPU 的虚拟机，请使用以下命令安装 JAX 和其他依赖项：

pip install -r vit_jax/requirements.txt

如果您连接的是带有 TPU 的虚拟机，请使用以下命令安装 JAX 和其他依赖项：

pip install -r vit_jax/requirements-tpu.txt

安装 Flaxformer，并按照此处链接的相应仓库中的说明进行操作。

无论是 GPU 还是 TPU，都请通过以下命令检查 JAX 是否能够连接到已挂载的加速器：

python -c 'import jax; print(jax.devices())'

最后，执行“微调模型”一节中提到的其中一个命令。

Bibtex

@article{dosovitskiy2020vit,
  title={一张图胜过 16×16 个词：大规模图像识别中的 Transformer},
  author={多索维茨基、卢卡斯·贝耶尔、亚历山大·科列斯尼科夫、迪尔克·魏森博恩、夏晓华·翟、托马斯·昂特尔蒂纳、莫斯塔法·德赫加尼、马蒂亚斯·明德勒、格奥尔格·海戈尔德、西尔万·盖利、雅各布·乌斯科雷特、尼尔·豪尔斯比},
  journal={ICLR},
  year={2021}
}

@article{tolstikhin2021mixer,
  title={MLP-Mixer：一种用于视觉任务的全 MLP 架构},
  author={伊利亚·托尔斯提欣、尼尔·豪尔斯比、亚历山大·科列斯尼科夫、卢卡斯·贝耶尔、夏晓华·翟、托马斯·昂特尔蒂纳、杰西卡·容、安德烈亚斯·施泰纳、丹尼尔·凯瑟斯、雅各布·乌斯科雷特、马里奥·卢奇奇、阿列克谢·多索维茨基},
  journal={arXiv 预印本 arXiv:2105.01601},
  year={2021}
}

@article{steiner2021augreg,
  title={如何训练你的 ViT？视觉 Transformer 中的数据、增强与正则化},
  author={安德烈亚斯·施泰纳、亚历山大·科列斯尼科夫、夏晓华·翟、罗斯·怀特曼、雅各布·乌斯科雷特、卢卡斯·贝耶尔},
  journal={arXiv 预印本 arXiv:2106.10270},
  year={2021}
}

@article{chen2021outperform,
  title={当 Vision Transformer 在无需预训练或强大数据增强的情况下超越 ResNet 时},
  author={陈翔宁、许祖辉、龚博清},
  journal={arXiv 预印本 arXiv:2106.01548},
  year={2021},
}

@article{zhuang2022gsam,
  title={代理差距最小化改进了尖锐度感知训练},
  author={庄俊堂、龚博清、袁良哲、崔寅、哈特维格·亚当、妮查·德沃内克、塞卡尔·塔蒂孔达、詹姆斯·邓肯、刘婷},
  journal={ICLR},
  year={2022},
}

@article{zhai2022lit,
  title={LiT：基于锁定图像文本调优的零样本迁移学习},
  author={夏晓华·翟、王肖、巴西尔·穆斯塔法、安德烈亚斯·施泰纳、丹尼尔·凯瑟斯、亚历山大·科列斯尼科夫、卢卡斯·贝耶尔},
  journal={CVPR},
  year={2022}
}

更改记录

按时间倒序排列：

• 2022-08-18：新增了在图像侧未使用线性头（LiT_B16B：768）且训练了6万步的LiT-B16B_2模型，其性能优于LiT_B16B。

• 2022-06-09：新增了在ImageNet数据集上使用GSAM从头开始训练的ViT和Mixer模型，未采用强数据增强。这些ViT模型的性能优于使用AdamW优化器、原始SAM算法或强数据增强训练的同规模模型。

• 2022-04-14：新增了LiT模型及其Colab笔记本。

• 2021-07-29：新增了ViT-B/8 AugReg模型（3个上游检查点及分辨率=224的适配版本）。

• 2021-07-02：新增了论文《当视觉Transformer超越ResNet时...》。

• 2021-07-02：新增了使用SAM（Sharpness-Aware Minimization，尖锐度感知最小化）优化的ViT和MLP-Mixer检查点。

• 2021-06-20：新增了论文《如何训练你的ViT？...》，并提供了一个新的Colab笔记本，用于探索论文中提到的5万多份预训练和微调检查点。

• 2021-06-18：本仓库已重写，采用Flax Linen API和ml_collections.ConfigDict进行配置管理。

• 2021-05-19：随着《如何训练你的ViT？...》论文的发表，我们新增了5万多份在ImageNet和ImageNet-21k数据集上预训练、并经过不同程度数据增强和模型正则化的ViT及混合模型，这些模型还在ImageNet、Pets37、Kitti-distance、CIFAR-100和Resisc45数据集上进行了微调。请查看vit_jax_augreg.ipynb来浏览这一丰富的模型资源！例如，您可以通过该Colab获取论文中表3 i21k_300列中推荐的预训练和微调检查点文件名。

• 2020-12-01：新增了R50+ViT-B/16混合模型（在ResNet-50骨干网络之上叠加ViT-B/16）。当在imagenet21k数据集上预训练时，该模型仅需不到一半的计算资源即可达到与L/16模型相近的性能。需要注意的是，“R50”在B/16变体中有所调整：原始ResNet-50包含[3,4,6,3]层，每层都会使图像分辨率降低两倍。结合ResNet的stem部分，最终分辨率将被缩小32倍，即使使用(1,1)的patch size，也无法实现ViT-B/16变体。因此，我们在R50+B/16变体中采用了[3,4,9]层的配置。

• 2020-11-09：新增了ViT-L/16模型。

• 2020-10-29：新增了在ImageNet-21k数据集上预训练、随后在ImageNet数据集上以224x224分辨率（而非默认的384x384）微调的ViT-B/16和ViT-L/16模型。这些模型名称后均带有“-224”后缀。预计它们的top-1准确率分别为81.2%和82.7%。

免责声明

开源发布由Andreas Steiner整理。

注：本仓库基于google-research/big_transfer分叉并修改而来。

这并非谷歌官方产品。

Vision Transformer (ViT) 快速上手指南

本指南基于 Google Research 开源的 vision_transformer 项目，帮助开发者快速在 JAX/Flax 环境下部署和微调 ViT 及 MLP-Mixer 模型。

环境准备

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

• 操作系统: Linux (推荐) 或 macOS。Windows 用户建议使用 WSL2 或 Docker。
• Python 版本: >= 3.10
• 硬件加速:
  • GPU: 支持 CUDA 的 NVIDIA 显卡（需安装对应的 CUDA 驱动）。
  • TPU: 可选，需配置 TPU 环境（如 Google Cloud TPU 或 Colab TPU）。
• 依赖框架: 本项目基于 JAX 和 Flax 构建。

注意：国内开发者若遇到 PyPI 下载缓慢问题，建议在 pip 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源。

安装步骤

1. 安装基础依赖

根据您使用的硬件类型（GPU 或 TPU），选择对应的依赖文件进行安装。

如果您使用 GPU：

pip install -r vit_jax/requirements.txt

如果您使用 TPU：

pip install -r vit_jax/requirements-tpu.txt

提示：如果上述命令因网络原因失败，请结合国内镜像源使用，例如： pip install -r vit_jax/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 Flaxformer

该项目依赖 Flaxformer 库，请参照其官方仓库说明进行安装：

pip install flaxformer

(注：如遇版本兼容问题，请参考 Flaxformer GitHub 获取最新安装指令)

基本使用

安装完成后，您可以直接加载预训练模型并在自定义数据集上进行微调（Fine-tuning）。所有模型共享统一的命令行接口。

场景一：微调 ViT 模型

以下示例展示如何将预训练在 ImageNet-21k 上的 ViT-B/16 模型微调到 CIFAR-10 数据集。代码会自动从 Google Cloud Storage (GCS) 拉取模型权重，无需手动下载。

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 \
    --config.pretrained_dir='gs://vit_models/imagenet21k'

场景二：微调 MLP-Mixer 模型

以下示例展示如何微调 Mixer-B/16 模型到 CIFAR-10 数据集：

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/mixer_base16_cifar10.py \
    --config.pretrained_dir='gs://mixer_models/imagenet21k'

场景三：使用增强正则化（AugReg）模型

如果您希望使用论文《How to train your ViT?》中提到的性能更强的 AugReg 预训练模型（例如在 Oxford-IIIT Pet 数据集上微调）：

python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \
    --config=$(pwd)/vit_jax/configs/augreg.py:R_Ti_16 \
    --config.dataset=oxford_iiit_pet \
    --config.base_lr=0.01

显存优化提示

如果在运行过程中遇到 Out Of Memory (OOM) 错误，可以通过调整以下参数来降低显存占用：

1. 增加梯度累积步数（减少每次反向传播的频率）： 添加参数 --config.accum_steps=8
2. 减小 Batch Size（需同步降低学习率）： 添加参数 --config.batch=256 并相应调整 --config.base_lr
3. 减小 Shuffle Buffer（解决主机内存不足）： 添加参数 --config.shuffle_buffer=10000

查看可用数据集与参数

要查看所有支持的命令行标志和数据集配置，可运行：

python3 -m vit_jax.train --help