懂AI
big_vision

big_vision

基于Jax/Flax的大规模视觉模型训练框架

Big Vision是一个用于训练大规模视觉模型的开源代码库。它基于Jax/Flax构建，支持在Cloud TPU VM和GPU上运行。该项目采用tf.data和TensorFlow Datasets实现高效的数据处理，可无缝扩展至2048个TPU核心的分布式环境。Big Vision涵盖了视觉Transformer、多模态学习、知识蒸馏等多个研究方向，为大规模视觉实验提供了可靠的基础。

big vision深度学习计算机视觉模型训练JaxGithub开源项目
GitHub论文

大视野

这个代码库旨在使用Cloud TPU VM或GPU机器训练大规模视觉模型。它基于Jax/Flax库，并使用tf.data和TensorFlow Datasets来实现可扩展和可复现的输入流水线。

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

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

big_vision旨在支持Google的研究项目。除非事先批准（请先在issue中询问），否则我们不太可能处理功能请求或接受外部贡献。对于一个仅用于迁移的、得到良好支持的代码库，请参见vision_transformer。

请注意，big_vision是一个相当动态的代码库，虽然我们打算始终保持核心代码完全功能，但我们不能保证及时更新位于.../proj/...子文件夹中的项目特定代码。然而，我们提供了一个表格，列出了已知特定项目工作的最后提交。

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

架构研究

多模态研究

训练

其他

代码库高层组织和原则简述

主要入口点是训练器模块，它通常完成所有与创建模型和优化器、加载数据、检查点和在循环中训练/评估模型相关的样板代码。我们在根文件夹中提供了规范的训练器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配置参数减少了默认批量大小和周期数):

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数据的说明。

具体来说,可以使用以下命令在TPU机器上的~/tensorflow_datasets下生成评估期间使用的七个TFDS数据集:

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上运行训练脚本

要在大型数据集(如imagenet2012)上训练您自己的big_vision模型(准备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-only机器的名称>
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

接下来,通过gcloud compute ssh $NAME_CPU_HOST --zone=$ZONEssh到机器,并按照说明格式化和挂载磁盘。 假设它被挂载到/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},
}

项目特定提交

特定项目代码预期工作的最后已知提交。核心代码和配置预期在最新版本中工作。

项目提交
UViMhttps://github.com/google-research/big_vision/commit/21bd6ebe253f070f584d8b777ad76f4abce51bef
image_texthttps://github.com/google-research/big_vision/commit/8921d5141504390a8a4f7b2dacb3b3c042237290
distillhttps://github.com/google-research/big_vision/commit/2f3f493af048dbfd97555ff6060f31a0e686f17f
GSAM进行中
CLIPPOhttps://github.com/google-research/big_vision/commit/fd2d3bd2efc9d89ea959f16cd2f58ae8a495cd44
CapPahttps://github.com/google-research/big_vision/commit/7ace659452dee4b68547575352c022a2eef587a5
GIVThttps://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)均在Apache2许可下发布。 有关完整的许可证文本,请参阅LICENSE文件。

编辑推荐精选

讯飞智文

讯飞智文

一键生成PPT和Word，让学习生活更轻松

讯飞智文是一个利用 AI 技术的项目，能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定，还是学生群体的职业生涯规划、实习避坑指南，亦或是活动策划、旅游攻略等内容，它都能提供支持，帮助用户精准表达，轻松呈现各种信息。

AI办公办公工具AI工具讯飞智文AI在线生成PPTAI撰写助手多语种文档生成AI自动配图热门
讯飞星火

讯飞星火

深度推理能力全新升级，全面对标OpenAI o1

科大讯飞的星火大模型，支持语言理解、知识问答和文本创作等多功能，适用于多种文件和业务场景，提升办公和日常生活的效率。讯飞星火是一个提供丰富智能服务的平台，涵盖科技资讯、图像创作、写作辅助、编程解答、科研文献解读等功能，能为不同需求的用户提供便捷高效的帮助，助力用户轻松获取信息、解决问题，满足多样化使用场景。

热门AI开发模型训练AI工具讯飞星火大模型智能问答内容创作多语种支持智慧生活
Spark-TTS

Spark-TTS

一种基于大语言模型的高效单流解耦语音令牌文本到语音合成模型

Spark-TTS 是一个基于 PyTorch 的开源文本到语音合成项目，由多个知名机构联合参与。该项目提供了高效的 LLM（大语言模型）驱动的语音合成方案，支持语音克隆和语音创建功能，可通过命令行界面（CLI）和 Web UI 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数，生成高质量的语音。该项目适用于多种场景，如有声读物制作、智能语音助手开发等。

Trae

Trae

字节跳动发布的AI编程神器IDE

Trae是一种自适应的集成开发环境（IDE），通过自动化和多元协作改变开发流程。利用Trae，团队能够更快速、精确地编写和部署代码，从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能，是提升开发效率的理想工具。

AI工具TraeAI IDE协作生产力转型热门
咔片PPT

咔片PPT

AI助力，做PPT更简单！

咔片是一款轻量化在线演示设计工具，借助 AI 技术，实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT，提供海量模板、智能美化、素材替换等功能，适用于销售、教师、学生等各类人群，能高效制作出高品质 PPT，满足不同场景演示需求。

讯飞绘文

讯飞绘文

选题、配图、成文，一站式创作，让内容运营更高效

讯飞绘文，一个AI集成平台，支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容，加速品牌传播，提升内容营销效果。

热门AI辅助写作AI工具讯飞绘文内容运营AI创作个性化文章多平台分发AI助手
材料星

材料星

专业的AI公文写作平台，公文写作神器

AI 材料星，专业的 AI 公文写作辅助平台，为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能，支持 30 + 文稿类型生成，助力快速完成领导讲话、工作总结、述职报告等材料，提升办公效率，是体制打工人的得力写作神器。

openai-agents-python

openai-agents-python

OpenAI Agents SDK，助力开发者便捷使用 OpenAI 相关功能。

openai-agents-python 是 OpenAI 推出的一款强大 Python SDK，它为开发者提供了与 OpenAI 模型交互的高效工具，支持工具调用、结果处理、追踪等功能，涵盖多种应用场景，如研究助手、财务研究等，能显著提升开发效率，让开发者更轻松地利用 OpenAI 的技术优势。

Hunyuan3D-2

Hunyuan3D-2

高分辨率纹理 3D 资产生成

Hunyuan3D-2 是腾讯开发的用于 3D 资产生成的强大工具，支持从文本描述、单张图片或多视角图片生成 3D 模型，具备快速形状生成能力，可生成带纹理的高质量 3D 模型，适用于多个领域，为 3D 创作提供了高效解决方案。

3FS

3FS

一个具备存储、管理和客户端操作等多种功能的分布式文件系统相关项目。

3FS 是一个功能强大的分布式文件系统项目，涵盖了存储引擎、元数据管理、客户端工具等多个模块。它支持多种文件操作，如创建文件和目录、设置布局等，同时具备高效的事件循环、节点选择和协程池管理等特性。适用于需要大规模数据存储和管理的场景，能够提高系统的性能和可靠性，是分布式存储领域的优质解决方案。

下拉加载更多