big_vision: Train and Evaluate Large-Scale Vision Models

Overview

big_vision is a research codebase for building, training, and evaluating large-scale vision models. It combines JAX and Flax with TensorFlow data pipelines, and is designed to scale from a single TPU core to distributed TPU setups, as well as GPU machines.

It is most useful for researchers who need training infrastructure and examples connected to model work such as ViT, MLP-Mixer, LiT, and SigLIP. The repository is research-oriented rather than a supported general-purpose product, and its maintainers say external contributions and feature requests are unlikely to be accepted without prior approval.

Key Features

• Trainer, model, evaluator, and preprocessing components organized for reuse across experiments.
• Configuration-driven training and evaluation, with runnable command examples in configuration files.
• Support for GPU machines and Google Cloud TPU VMs, including distributed training.
• Checkpointing that can resume interrupted training jobs when the same work directory is provided.
• Flexible parameter and model sharding, including FSDP configuration for TPU fine-tuning.
• Code and configurations associated with vision and multimodal research projects, including ViT, MLP-Mixer, LiT, and SigLIP.
• TFDS-based data preparation for reproducible access to datasets.

Use Cases

• Vision researchers can use the included trainers and configurations as a starting point for large-scale model experiments.
• Teams with Cloud TPU access can run distributed training and fine-tuning without assembling all the training infrastructure from scratch.
• Practitioners can fine-tune supported pretrained vision models on downstream datasets using the transfer configuration.
• Researchers exploring image-text learning can examine or run the repository’s contrastive training examples, including SigLIP configurations.

Getting Started

Clone the repository and install its listed dependencies:

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

Training also requires a compatible JAX installation, prepared dataset, and GPU or TPU resources. See the README for setup and run instructions.

Considerations

• Training large models requires suitable accelerator hardware and can involve substantial data and infrastructure setup.
• JAX installation depends on the machine’s CUDA and cuDNN setup when using GPUs.
• Some datasets require manual downloads and preparation.
• The maintainers describe the codebase as dynamic. Core code is expected to work at the repository head, but project-specific code under proj may rely on older commits.
• The repository is primarily intended to support Google Research projects. External feature requests and contributions are not generally expected to receive support unless pre-approved.