大型视觉模型

该代码库旨在使用Cloud TPU VMs或 GPU 机器来训练大规模视觉模型。它基于Jax/Flax库，并利用tf.data和TensorFlow Datasets构建可扩展且可复现的数据流水线。

开源此代码库主要有两个目的：

1. 发布在此代码库中开发的研究项目代码（见下方列表）。
2. 为在 GPU 机器和 Google Cloud TPUs 上运行大规模视觉实验提供一个强有力的起点，能够从单个 TPU 核心无缝扩展到最多包含 2048 个 TPU 核心的分布式设置，开箱即用。

big_vision 的目标是支持 Google 内部的研究项目。我们通常不会处理功能请求或接受外部贡献，除非事先获得批准（请先在 issue 中咨询）。对于一个得到良好支持的纯迁移代码库，也可以参考vision_transformer。

请注意，big_vision 是一个非常动态的代码库，尽管我们致力于始终保持核心代码的完全可用性，但我们无法保证位于 .../proj/... 子文件夹中的项目特定代码能够及时更新。不过，我们提供了一个表格，列出了已知特定项目可以正常工作的最后提交记录。

以下研究项目最初是在 big_vision 代码库中开展的：

架构研究

• 一张图胜过 16x16 个词：大规模图像识别中的 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 的规模效应，作者：Xiaohua Zhai*、Alexander Kolesnikov*、Neil Houlsby 和 Lucas Beyer*
  资源：配置文件。
• 如何训练你的 ViT？视觉 Transformer 中的数据、增强与正则化，作者：Andreas Steiner*、Alexander Kolesnikov*、Xiaohua Zhai*、Ross Wightman、Jakob Uszkoreit 和 Lucas Beyer*
• MLP-Mixer：一种全 MLP 的视觉架构，作者：Ilya Tolstikhin*、Neil Houlsby*、Alexander Kolesnikov*、Lucas Beyer*、Xiaohua Zhai、Thomas Unterthiner、Jessica Yung、Andreas Steiner、Daniel Keysers、Jakob Uszkoreit、Mario Lucic 和 Alexey Dosovitskiy
  资源：配置文件。
• ImageNet-1k 数据集上更好的纯 ViT 基线，作者：Lucas Beyer、Xiaohua Zhai、Alexander Kolesnikov
  资源：配置文件。
• UViM：一种基于学习引导码的统一视觉建模方法，作者：Alexander Kolesnikov^、André Susano Pinto^、Lucas Beyer*、Xiaohua Zhai*、Jeremiah Harmsen* 和 Neil Houlsby*
  资源：README 文件、配置文件、Colab 笔记本。
• FlexiViT：一种适用于所有补丁大小的单一模型，作者：Lucas Beyer*、Pavel Izmailov*、Alexander Kolesnikov*、Mathilde Caron*、Simon Kornblith*、Xiaohua Zhai*、Matthias Minderer*、Michael Tschannen*、Ibrahim Alabdulmohsin* 和 Filip Pavetic*
  资源：README 文件、配置文件。
• 双层 PatchNorm，作者：Manoj Kumar、Mostafa Dehghani 和 Neil Houlsby。
• 让 ViT 达到最佳状态：计算最优模型设计的规模法则，作者：Ibrahim Alabdulmohsin*、Xiaohua Zhai*、Alexander Kolesnikov 和 Lucas Beyer*。
• （部分）将视觉 Transformer 扩展至 220 亿参数，作者：Mostafa Dehghani*、Josip Djolonga*、Basil Mustafa*、Piotr Padlewski*、Jonathan Heek*、众多中间作者 和 Neil Houlsby*。
• （部分）有限标量量化：简化版 VQ-VAE，作者：Fabian Mentzer、David Minnen、Eirikur Agustsson 和 Michael Tschannen。
• GIVT：生成式无限词汇表 Transformer，作者：Michael Tschannen、Cian Eastwood 和 Fabian Mentzer。
  资源：README 文件、配置文件、Colab 笔记本。
• 基于掩码扩散的统一自编码，作者：Philippe Hansen-Estruch、Sriram Vishwanath、Amy Zhang 和 Manan Tomar。
• Jet：一种现代基于 Transformer 的归一化流，作者：Alexander Kolesnikov*、André Susano Pinto*、Michael Tschannen*、配置文件。
• JetFormer：一种自回归生成模型，用于处理原始图像和文本，作者：Michael Tschannen*、André Susano Pinto* 和 Alexander Kolesnikov*。配置文件。

多模态研究

• LiT：锁定图像的文本调优实现零样本迁移，作者： Xiaohua Zhai*、Xiao Wang*、Basil Mustafa*、Andreas Steiner*、Daniel Keysers、Alexander Kolesnikov 和 Lucas Beyer*
  资源：训练器、配置文件、Colab 笔记本。
• CLIPPO：仅从像素进行图像与语言理解，作者： Michael Tschannen、Basil Mustafa、Neil Houlsby
  资源：README、配置文件、Colab 笔记本。
• 用于语言图像预训练的 Sigmoid 损失，作者： Xiaohua Zhai*、Basil Mustafa、Alexander Kolesnikov、Lucas Beyer*
  资源：Colab 笔记本及模型，代码待补充。
• 计算机视觉多任务学习中自回归解码器的研究，作者： Lucas Beyer*、Bo Wan*、Gagan Madan*、Filip Pavetic*、Andreas Steiner*、Alexander Kolesnikov、André Susano Pinto、Emanuele Bugliarello、Xiao Wang、Qihang Yu、Liang-Chieh Chen、Xiaohua Zhai*。
• 图像字幕生成器同样是可扩展的视觉学习者，作者： Michael Tschannen*、Manoj Kumar*、Andreas Steiner*、Xiaohua Zhai、Neil Houlsby、Lucas Beyer*。
  资源：README、配置文件、模型。
• 三塔：基于预训练图像模型的灵活对比学习，作者：Jannik Kossen、Mark Collier、Basil Mustafa、Xiao Wang、Xiaohua Zhai、Lucas Beyer、Andreas Steiner、Jesse Berent、Rodolphe Jenatton、Efi Kokiopoulou。
• （部分）PaLI：联合缩放的多语言语言-图像模型，作者：Xi Chen、Xiao Wang、Soravit Changpinyo、哇，中间作者真多、Anelia Angelova、Xiaohua Zhai、Neil Houlsby、Radu Soricut。
• （部分）PaLI-3 视觉语言模型：更小、更快、更强，作者：Xi Chen、Xiao Wang、Lucas Beyer、Alexander Kolesnikov、Jialin Wu、Paul Voigtlaender、Basil Mustafa、Sebastian Goodman、Ibrahim Alabdulmohsin、Piotr Padlewski、Daniel Salz、Xi Xiong、Daniel Vlasic、Filip Pavetic、Keran Rong、Tianli Yu、Daniel Keysers、Xiaohua Zhai、Radu Soricut。
• LocCa，作者： Bo Wan、Michael Tschannen、Yongqin Xian、Filip Pavetic、Ibrahim Alabdulmohsin、Xiao Wang、André Susano Pinto、Andreas Steiner、Lucas Beyer、Xiaohua Zhai。
• PaliGemma、 PaliGemma 2，作者：哇，作者真多。
  资源：README、 模型、 迁移配置、 数据集、 CountBenchQA。
• SigLIP 2：具有改进语义理解、定位能力和密集特征的多语言视觉-语言编码器，作者：哇，作者真多。
  资源：README（含检查点）、Colab 笔记本。

训练

• 知识蒸馏：好老师要有耐心和一致性，作者： Lucas Beyer*、Xiaohua Zhai*、Amélie Royer*、Larisa Markeeva*、Rohan Anil、以及 Alexander Kolesnikov*
  资源：README、训练器、Colab 笔记本。
• 尖锐度感知最小化：高效提升泛化能力的方法，作者： Pierre Foret、Ariel Kleiner、Hossein Mobahi、Behnam Neyshabur。
• 代理差距最小化改进尖锐度感知训练，作者：Juntang Zhuang、Boqing Gong、Liangzhe Yuan、Yin Cui、Hartwig Adam、Nicha Dvornek、Sekhar Tatikonda、James Duncan 和 Ting Liu
  资源：训练器、配置文件、复现结果。
• 使用任务奖励微调计算机视觉模型，作者： André Susano Pinto*、Alexander Kolesnikov*、Yuge Shi、Lucas Beyer、Xiaohua Zhai。
• （部分）VeLO：通过扩大规模训练多功能学习优化器，作者： Luke Metz、James Harrison、C. Daniel Freeman、Amil Merchant、Lucas Beyer、James Bradbury、Naman Agrawal、Ben Poole、Igor Mordatch、Adam Roberts、Jascha Sohl-Dickstein。

杂项

• 我们已经用完 ImageNet 了吗？，作者： Lucas Beyer*、Olivier J. Hénaff*、Alexander Kolesnikov*、Xiaohua Zhai*、Aäron van den Oord*。
• 无过滤：对比式视觉-语言模型中的文化与社会经济多样性，作者： Angéline Pouget、Lucas Beyer、Emanuele Bugliarello、Xiao Wang、Andreas Peter Steiner、Xiaohua Zhai、Ibrahim Alabdulmohsin。

代码库的高层次组织与原则概览

主入口是一个训练器模块，它通常负责创建模型和优化器、加载数据、保存检查点以及在循环中进行训练和评估等所有样板代码。我们在根目录下提供了标准的训练器 train.py。通常，big_vision 中的各个项目会基于此训练器进行分支并定制。

所有模型、评估器和预处理操作都位于相应的子目录中，并且经常可以在不同项目之间复用。我们鼓励这些目录内的 API 兼容，以促进复用性，但这并不是强制性的，因为各个项目可能需要引入自己的自定义 API。

我们有一个强大的配置系统，配置文件存放在 configs/ 目录中。自定义的训练器和模块可以直接扩展或修改配置选项。

项目特定的代码位于 .../proj/... 命名空间中。由于各种原因，项目特定的代码并不总是能够与核心 big_vision 库保持同步。下面列出了每个项目在预期可以正常工作的 最新已知提交。

训练任务对中断具有鲁棒性，可以从上次保存的检查点无缝恢复（前提是用户提供了正确的 --workdir 路径）。

每个配置文件的顶部都包含一个注释，其中有一段用于运行该配置的 COMMAND 代码片段，以及对预期运行时间和结果的一些提示。更多详细信息请见下文，但一般来说，在 GPU 机器上运行时只需调用 python -m COMMAND；而在 TPU 上运行时（包括多主机场景），则需要执行以下命令：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all
  --command "bash big_vision/run_tpu.sh COMMAND"

有关如何在 GPU 机器或 Google Cloud TPU 上运行 big_vision 代码的详细说明，请参阅下方内容。

默认情况下，我们会写入检查点和日志文件。日志文件是以 JSON 对象组成的列表，我们提供了一个简短而直接的 示例 Colab，用于读取和显示日志及检查点。

当前与未来的内容

首次发布包含了在 Cloud TPU VM 上大规模预训练、迁移和评估分类模型的核心部分。

此后，我们又添加了以下关键功能和项目：

• 类似 LiT 和 CLIP 的对比图像-文本模型训练与评估。
• 稳健且一致的蒸馏技术。
• ViT 的规模扩展。
• MLP-Mixer。
• UViM。

我们计划在近期发布的功能和项目（不分先后顺序）包括：

• TFDS 中的 ImageNet-21k 数据集。
• 加载我们在论文中使用的其他公开模型（NFNet、MoCov3、DINO）。
• 内存高效的 Polyak 平滑实现。
• 高级 JAX 计算与内存分析工具。目前我们使用的是内部工具，但未来可能会增加对公开可用工具的支持。

我们将继续在此处发布在 big_vision 框架内开发的未来研究成果代码。

非内容部分

以下内容存在于该代码库的内部版本中，目前暂无对外发布的计划：

• 针对质量和速度的常规回归测试，它们高度依赖于内部基础设施。
• 实验的高级日志记录、监控和可视化工具，这也严重依赖于内部基础设施。不过，我们对此持开放态度，未来可能会加入一些功能，尤其是如果能以独立的方式实现的话。
• 尚未发表的正在进行中的研究项目。

GPU 设置

我们首先讨论如何在本地 GPU 机器上设置和运行 big_vision，然后再介绍 Cloud TPU 的设置方法。需要注意的是，本地 GPU 设置中的数据准备步骤在很大程度上也可以用于 Cloud TPU 设置。虽然为了简洁起见，相关说明中省略了这一步骤，但我们强烈建议在安装 Python 依赖时使用 虚拟环境。

安装 Python 包

第一步是克隆 big_vision 仓库并安装相关的 Python 依赖：

git clone https://github.com/google-research/big_vision
cd big_vision/
pip3 install --upgrade pip
pip3 install -r big_vision/requirements.txt

最新版本的 jax 库可以通过以下命令获取：

pip3 install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

根据您机器上安装的 CUDA 和 cuDNN 库版本，您可能需要选择不同的 jax 包。更多信息请参考 官方 JAX 文档。

准备 tfds 数据

为了统一且可复现地访问标准数据集，我们选择了 tensorflow_datasets (tfds) 库。使用该库时，需要先下载并预处理每个数据集，然后将其存储在硬盘上（或者，如果您使用 Google Cloud，则最好存储在 GCP 存储桶中）。

许多数据集在首次使用时可以自动下载和预处理。尽管如此，我们特意禁用了这一功能，建议在首次运行之前单独完成数据准备步骤。这样可以在出现问题时更方便地进行调试，而且某些数据集，如 imagenet2012，需要手动下载数据。

大多数数据集，例如 cifar100、oxford_iiit_pet 或 imagenet_v2，都可以通过运行以下命令完全自动地下载和准备：

cd big_vision/
python3 -m big_vision.tools.download_tfds_datasets cifar100 oxford_iiit_pet imagenet_v2

完整的数据集列表可在 此链接 查看。

然而，有些数据集，如 imagenet2012 或 imagenet2012_real，需要手动下载数据并放置到 $TFDS_DATA_DIR/downloads/manual/ 目录中，该目录的默认路径为 ~/tensorflow_datasets/downloads/manual/。例如，对于 imagenet2012 和 imagenet2012_real，需要将官方的 ILSVRC2012_img_train.tar 和 ILSVRC2012_img_val.tar 文件放入该目录，然后运行：

python3 -m big_vision.tools.download_tfds_datasets imagenet2012 imagenet2012_real

（这可能需要大约 1 小时。）

如果您使用 Google Cloud，特别是 TPU，可以将预处理后的数据（存储在 $TFDS_DATA_DIR 中）上传到 Google Cloud 存储桶，并在任何 TPU 虚拟机上使用该存储桶来访问数据。

在 GPU 机器上运行

最后，在安装完所有 Python 依赖并准备好 tfds 数据后，用户可以使用自己选择的配置来运行任务。例如，要在 ImageNet 数据上训练 ViT-S/16 模型，应运行以下命令：

python3 -m big_vision.train --config big_vision/configs/vit_s16_i1k.py --workdir workdirs/`date '+%m-%d_%H%M'`

或者要训练 MLP-Mixer-B/16，可以运行（注意 gpu8 配置参数会减少默认的批大小和 epoch 数）：

python3 -m big_vision.train --config big_vision/configs/mlp_mixer_i1k.py:gpu8 --workdir workdirs/`date '+%m-%d_%H%M'`

Cloud TPU VM 设置

创建 TPU VM

要创建一台拥有 8 个 TPU 核心的单机，请按照以下 Cloud TPU JAX 文档操作： https://cloud.google.com/tpu/docs/run-calculation-jax

为了支持大规模视觉研究，建议使用多台主机以获得更多核心。下面我们将提供如何实现这一目标的说明。

首先，创建一些可重复使用的变量：

export NAME=<TPU 部署名称，例如 my-tpu-machine>
export ZONE=<GCP 地理区域，例如 europe-west4-a>
export GS_BUCKET_NAME=<存储桶名称，例如 my_bucket>

以下命令行将创建一个拥有 32 个核心、4 台主机的 TPU VM：

gcloud compute tpus tpu-vm create $NAME --zone $ZONE --accelerator-type v3-32 --version tpu-ubuntu2204-base

在 TPU VM 上安装 big_vision

获取 big_vision 仓库，将其复制到所有 TPU VM 主机，并安装依赖项。

git clone https://github.com/google-research/big_vision
gcloud compute tpus tpu-vm scp --recurse big_vision/big_vision $NAME: --zone=$ZONE --worker=all
gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "bash big_vision/run_tpu.sh"

下载并准备 TFDS 数据集

我们建议先按照上述方法在本地准备 tfds 数据，然后再将其上传到 Google Cloud 存储桶。不过，如果您愿意，也可以使用 TPU 机器自动准备那些无需手动下载的数据集，具体步骤如下。请注意，TPU 机器仅有 100 GB 的磁盘空间，且多主机 TPU 分片不允许以写入模式挂载外部磁盘，因此以下说明可能不适用于大型数据集的准备。作为另一种替代方案，我们还提供了 如何在仅 CPU 的 GCP 机器上准备 tfds 数据 的说明。

具体来说，评估过程中使用的七个 TFDS 数据集将在 TPU 机器上的 ~/tensorflow_datasets 目录下通过以下命令生成：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=0 --command "TFDS_DATA_DIR=~/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.tools.download_tfds_datasets cifar10 cifar100 oxford_iiit_pet oxford_flowers102 cars196 dtd uc_merced"

随后，您可以将这些数据集复制到 GS 存储桶，以便所有 TPU 工作节点都能访问。

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=0 --command "rm -r ~/tensorflow_datasets/downloads && gsutil cp -r ~/tensorflow_datasets gs://$GS_BUCKET_NAME"

如果您希望集成其他公共或自定义数据集，例如 imagenet2012，请遵循 官方指南。

预训练模型

有关预训练模型的完整列表，请查看与模型代码位于同一模块中的 load 函数。关于如何使用这些模型的示例配置，请参阅 configs/transfer.py。

在 TPU VM 上运行迁移学习脚本

以下命令行将在 cifar10 数据集上微调一个预训练的 vit-i21k-augreg-b/32 模型：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/transfer.py:model=vit-i21k-augreg-b/32,dataset=cifar10,crop=resmall_crop --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'` --config.lr=0.03"

在 TPU VM 上运行训练脚本

要在一个大型数据集上训练您自己的 big_vision 模型，例如 imagenet2012（请先 准备 TFDS 数据集），可以运行以下命令行：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/bit_i1k.py  --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'`"

FSDP 训练

big_vision 支持灵活的参数和模型分片策略。目前，我们通过简单的配置更改即可支持流行的 FSDP 分片，详情请参见 此配置示例。例如，要对预训练的 ViT-L 模型进行 FSDP 微调，可以运行以下命令（可根据您的硬件调整批大小）：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/transfer.py:model=vit-i21k-augreg-l/16,dataset=oxford_iiit_pet,crop=resmall_crop,fsdp=True,batch_size=256 --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'` --config.lr=0.03"

使用 SigLIP 进行图像-文本联合训练

一个使用公开 coco 字幕数据的最小示例：

gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.trainers.proj.image_text.siglip --config big_vision/configs/proj/image_text/siglip_lit_coco.py --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%Y-%m-%d_%H%M'`"

有时有用的 gcloud 命令

• 销毁 TPU 机器：gcloud compute tpus tpu-vm delete $NAME --zone $ZONE
• 删除所有主机上的 big_vision 相关文件夹：gcloud compute tpus tpu-vm ssh $NAME --zone $ZONE --worker=all --command 'rm -rf ~/big_vision ~/bv_venv'

在独立的 GCP CPU 机器上准备 tfds 数据。

首先创建一台新机器和一块磁盘（您可以根据需要调整具体的机器类型和磁盘设置/容量）：

export NAME_CPU_HOST=<仅CPU机器的名称>
export NAME_DISK=<磁盘的名称>
gcloud compute instances create $NAME_CPU_HOST --machine-type c3-standard-22 --zone $ZONE --image-family ubuntu-2204-lts --image-project ubuntu-os-cloud
gcloud compute disks create $NAME_DISK --size 1000GB --zone $ZONE --type pd-balanced

现在将磁盘挂载到新创建的机器上：

gcloud compute instances attach-disk $NAME_CPU_HOST --disk $NAME_DISK --zone $ZONE

接下来，通过 ssh 登录到该机器：gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE，并 按照说明格式化和挂载磁盘。 假设它被挂载到了 /mnt/disks/tfds。

差不多完成了，现在克隆并设置 big_vision：

gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "git clone https://github.com/google-research/big_vision.git && cd big_vision && sh big_vision/run_tpu.sh"

最后，使用工具脚本准备数据集（例如 coco_captions），并将结果复制到您的 Google Cloud 存储桶中：

gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "cd big_vision && TFDS_DATA_DIR=/mnt/disks/tfds/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.tools.download_tfds_datasets coco_captions"
gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "rm -rf /mnt/disks/tfds/tensorflow_datasets/downloads && gsutil cp -r /mnt/disks/tfds/tensorflow_datasets gs://$GS_BUCKET_NAME"

ViT 基线

我们在名为 vit_s16_i1k.py 的配置文件中提供了一个经过良好调优的 ViT-S/16 基线。在 90 个训练周期内，它在 ImageNet 验证集上达到了 76.5% 的准确率，是研究 ViT 模型的一个强大而简单的起点。

请参阅我们的 arXiv 报告 以获取更多详细信息；如果这个基线对您的研究有所帮助，请考虑引用以下内容：

@article{vit_baseline,
  url = {https://arxiv.org/abs/2205.01580},
  author = {Beyer, Lucas and Zhai, Xiaohua and Kolesnikov, Alexander},
  title = {Better plain ViT baselines for ImageNet-1k},
  journal={arXiv preprint arXiv:2205.01580},
  year = {2022},
}

项目特定提交记录

以下是已知的最后一个提交记录，预计该项目代码在此版本下能够正常运行。核心代码和配置则应在最新版本下正常工作。

项目	提交
UViM	https://github.com/google-research/big_vision/commit/21bd6ebe253f070f584d8b777ad76f4abce51bef
image_text	https://github.com/google-research/big_vision/commit/8921d5141504390a8a4f7b2dacb3b3c042237290
distill	https://github.com/google-research/big_vision/commit/2f3f493af048dbfd97555ff6060f31a0e686f17f
GSAM	开发中
CLIPPO	https://github.com/google-research/big_vision/commit/fd2d3bd2efc9d89ea959f16cd2f58ae8a495cd44
CapPa	https://github.com/google-research/big_vision/commit/7ace659452dee4b68547575352c022a2eef587a5
GIVT	https://github.com/google-research/big_vision/commit/0cb70881dd33b3343b769347dc19793c4994b8cb

引用代码库

如果您发现此代码库对您的研究有所帮助，请考虑使用以下 BibTEX 格式来引用它：

@misc{big_vision,
  author = {Beyer, Lucas and Zhai, Xiaohua and Kolesnikov, Alexander},
  title = {Big Vision},
  year = {2022},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/google-research/big_vision}}
}

免责声明

本项目并非 Google 官方产品。

许可证

除非另有明确说明，否则 big_vision 代码库中的所有内容（包括模型和 Colab 笔记本）均采用 Apache 2.0 许可证发布。 完整的许可证文本请参阅 LICENSE 文件。

Big Vision 快速上手指南

big_vision 是 Google Research 开源的大规模视觉模型训练代码库，基于 Jax/Flax 构建，专为在 Cloud TPU VMs 或 GPU 机器上扩展训练而设计。它支持从单核到 2048 核 TPU 的无缝分布式训练，涵盖了 ViT、MLP-Mixer、SigLIP、PaliGemma 等前沿研究项目。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• 硬件:
  • GPU: 支持 CUDA 的 NVIDIA GPU (需安装对应驱动)
  • TPU: Google Cloud TPU VMs (推荐用于大规模实验)
• Python: 3.10 或更高版本

前置依赖

确保已安装以下基础工具：

• git
• python3-pip
• virtualenv (可选但推荐)

注意：本项目重度依赖 JAX 生态。若使用 GPU，请确保已正确安装 CUDA Toolkit 和 cuDNN。国内用户建议使用阿里云或清华源加速 Python 包下载。

安装步骤

1. 克隆仓库

   git clone https://github.com/google-research/big_vision.git
   cd big_vision
   

2. 创建虚拟环境并激活

   python3 -m venv venv
   source venv/bin/activate
   

3. 安装依赖

   方案 A：通用安装 (适用于大多数 GPU/CPU 环境) 先升级 pip，然后安装核心依赖。国内用户可临时切换镜像源：

   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade pip
   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -e .
   

   注：-e . 表示以编辑模式安装，方便修改代码。

   方案 B：针对 TPU 环境 (Google Cloud TPU VM) 如果在 TPU VM 上运行，通常预装了相关环境，只需安装项目特定依赖：

   pip install -e .
   

   若需特定版本的 JAX/Flax 以匹配 TPU 运行时，请参考 Cloud TPU JAX 安装指南。

4. 验证安装 尝试导入核心库确认无报错：

   python -c "import big_vision; print('Big Vision installed successfully')"
   

基本使用

big_vision 的核心入口是训练器脚本 train.py。它通过配置文件驱动模型架构、数据集、优化器和训练策略。

1. 最简单的运行示例

以下命令使用默认配置在本地（单 GPU 或 CPU）启动一个小型 ViT 模型的训练演示。该配置通常用于快速验证环境是否就绪。

python train.py --config big_vision/configs/vit_s16_i1k.py

说明：此命令会加载 ImageNet-1k 的小型 ViT 配置。首次运行会自动下载数据集（需联网），请确保存储空间充足。

2. 使用特定研究项目配置

big_vision 包含大量复现论文的配置。例如，运行 SigLIP (Sigmoid Loss for Language Image Pre-Training) 的演示：

# 运行 SigLIP 演示配置 (需根据实际文件路径调整)
python train.py --config big_vision/configs/proj/image_text/lit_coco.py

或者运行 MLP-Mixer 的训练：

python train.py --config big_vision/configs/mlp_mixer_i1k.py

3. 关键参数说明

在执行 train.py 时，常用参数如下：

• --config: 指定配置文件路径 (.py 文件)。
• --workdir: 指定模型检查点、日志和 TensorBoard 数据的保存目录 (默认为当前目录下的 workdir)。
• --mode: 运行模式，如 train (训练), eval (评估), write_results (输出结果)。

示例：指定工作目录进行训练

python train.py \
  --config big_vision/configs/vit_s16_i1k.py \
  --workdir /path/to/your/experiment_output

4. 在 Colab 中体验

对于不想配置本地环境的用户，许多项目提供了直接的 Colab 笔记本。例如 GIVT 或 PaliGemma 的 Demo：

• 访问仓库中的 big_vision/configs/proj/... 目录查找 .ipynb 文件链接。
• 点击 "Open in Colab" 即可在云端免费 GPU/TPU 上运行。

---

提示：big_vision 是一个动态的研究代码库，部分位于 proj/ 子文件夹中的特定项目代码可能随时间变化。建议在使用特定论文代码前，查阅该子文件夹下的 README.md 获取最新的 commit 哈希值或特定说明。