Big Vision:谷歌研究院的大规模视觉模型训练框架

big_vision

Big Vision:谷歌研究院的大规模视觉模型训练框架

Big Vision是谷歌研究院开源的一个用于训练大规模视觉模型的代码库。该项目旨在为研究人员提供一个强大的起点,以在GPU机器和Google Cloud TPU上进行大规模视觉实验。Big Vision支持多种先进的视觉模型架构,如Vision Transformer (ViT)、MLP-Mixer、SigLIP等,并提供了高效的训练和评估流程。

项目概述

Big Vision代码库的设计目标主要有两个:

发布在该代码库中开发的研究项目的代码。
为在GPU机器和Google Cloud TPU上运行大规模视觉实验提供一个强大的起点,可以无缝地从单个TPU核心扩展到多达2048个TPU核心的分布式设置。

该项目使用JAX/Flax库进行模型实现,使用tf.data和TensorFlow Datasets进行可扩展和可重现的输入流水线。Big Vision支持多种视觉和多模态模型,包括:

Vision Transformer (ViT)
MLP-Mixer
SigLIP (用于图像-文本预训练)
UViM (统一的视觉建模方法)
LiT (用于零样本迁移的锁定图像文本调优)
等等

主要特性

Big Vision具有以下主要特性:

支持多种先进的视觉模型架构
可在GPU和Cloud TPU上高效训练
使用tf.data和TensorFlow Datasets实现可扩展的输入流水线
支持大规模分布式训练
提供预训练模型和配置文件
实现了多种训练技巧,如知识蒸馏、锐度感知最小化等
支持多模态训练,如图像-文本对比学习

代码结构

Big Vision的代码结构清晰明了:

主入口点是trainer模块,通常负责创建模型和优化器、加载数据、检查点和训练/评估模型等所有样板代码。
所有模型、评估器和预处理操作都位于相应的子目录中,可以在不同项目之间重用。
配置文件位于configs/目录中。
项目特定的代码位于.../proj/...命名空间中。

使用方法

要使用Big Vision,首先需要设置Python环境并安装依赖项:

git clone https://github.com/google-research/big_vision
cd big_vision/
pip install -r big_vision/requirements.txt
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

然后,需要准备训练数据。Big Vision使用TensorFlow Datasets (tfds)来访问标准数据集。可以使用以下命令准备数据集:

python3 -m big_vision.tools.download_tfds_datasets cifar100 oxford_iiit_pet imagenet_v2

对于ImageNet等需要手动下载的数据集,需要先下载数据文件并放置在正确的目录中。

准备好环境和数据后,就可以开始训练模型了。例如,要在ImageNet上训练ViT-S/16模型,可以运行:

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

Big Vision还支持在Cloud TPU上进行大规模分布式训练。使用提供的脚本可以轻松设置TPU环境并启动训练任务。

预训练模型

Big Vision提供了多个预训练模型,可以直接用于迁移学习或进一步微调。例如,可以使用以下命令在CIFAR-10数据集上微调预训练的ViT模型:

python3 -m big_vision.train --config big_vision/configs/transfer.py:model=vit-i21k-augreg-b/32,dataset=cifar10,crop=resmall_crop --workdir workdir/finetune --config.lr=0.03

研究项目

Big Vision已被用于开发多个重要的计算机视觉研究项目,包括:

Vision Transformer (ViT)
MLP-Mixer
LiT (Locked-image Text Tuning)
SigLIP (Sigmoid Loss for Language Image Pre-Training)
UViM (A Unified Modeling Approach for Vision with Learned Guiding Codes)
FlexiViT (One Model for All Patch Sizes)
CLIPPO (Image-and-Language Understanding from Pixels Only)

这些项目涵盖了架构研究、多模态学习、训练方法改进等多个方面,推动了计算机视觉领域的发展。

结论

Big Vision为计算机视觉研究提供了一个强大而灵活的框架。它支持最新的模型架构,可在各种硬件上高效训练,并已被用于开发多个重要的研究项目。无论是想复现已发表的研究结果,还是开展新的视觉模型研究,Big Vision都是一个值得考虑的工具。

对于有兴趣使用或贡献Big Vision的研究人员,可以访问项目的GitHub仓库获取更多信息:https://github.com/google-research/big_vision。该项目仍在积极开发中,未来可能会添加更多功能和项目支持。

总的来说,Big Vision代表了谷歌研究院在推动大规模视觉模型研究方面的努力,为整个计算机视觉社区提供了宝贵的资源。随着视觉AI的不断发展,像Big Vision这样的开源工具将在推动技术进步和促进协作研究方面发挥重要作用。