google-research/vision_transformer

这是什么？

核心功能

• 01以 JAX 与 Flax 实现 ViT 和 MLP-Mixer 模型。
• 02提供在 ImageNet 和 ImageNet-21k 上预训练的检查点。
• 03通过命令行微调脚本，可接入 TensorFlow Datasets 中的数据集。
• 04提供 5 万余个采用不同数据增强设置的 ViT 与混合模型权重。
• 05附带用于模型探索、微调和 LiT 推理的 Colab 笔记本。

适用场景

• 在自定义图像分类数据集上微调预训练的 ViT 或 Mixer 模型。
• 研究 ViT 与 MLP-Mixer 架构及其训练方案的差异。
• 使用已发布的 LiT 模型进行零样本图文推理。

快速上手

1. 1

   确保已安装 Python 3.10 或更高版本

2. 2

   为 GPU 安装 JAX 及依赖

   pip install -r vit_jax/requirements.txt
3. 3

   在 CIFAR-10 上微调 ViT-B/16 模型（在 ImageNet-21k 上预训练）

   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'

使用前须知

• !微调需要 GPU 或 TPU，且 README 说明会使用所有可用的加速器。
• !Colab 上的 TPU 训练速度较慢，处理较大数据集时建议使用专用机器。
• !安装 JAX 和 Flaxformer 需要分别参考它们各自仓库的说明。

收录于专题

标签

常见问题

vision_transformer 是什么？

基于 JAX/Flax 的 Vision Transformer 与 MLP-Mixer 图像分类模型代码及预训练权重。 该仓库发布了 Vision Transformer（ViT）和 MLP-Mixer 两类图像识别模型的代码与预训练权重。ViT 将图像切分为固定大小的图块后送入 Transformer 编码器，MLP-Mixer 则完全由多层感知机构成。仓库同时对应 AugReg、LiT 零样本图文模型以及锐度感知训练等多篇论文的成果。

vision_transformer 可以免费使用吗？用的什么许可证？

vision_transformer 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。

vision_transformer 还在维护吗？

半年内有更新。vision_transformer 最近一次提交是在 2026年10月1日。

vision_transformer 适合用来做什么？

在自定义图像分类数据集上微调预训练的 ViT 或 Mixer 模型。 研究 ViT 与 MLP-Mixer 架构及其训练方案的差异。 使用已发布的 LiT 模型进行零样本图文推理。

vision_transformer 怎么安装和运行？

确保已安装 Python 3.10 或更高版本 为 GPU 安装 JAX 及依赖: pip install -r vit_jax/requirements.txt 在 CIFAR-10 上微调 ViT-B/16 模型（在 ImageNet-21k 上预训练）: 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'

在哪里获取 vision_transformer？

源代码和安装说明在 GitHub 上：https://github.com/google-research/vision_transformer