{"name":"big_vision: Google Research's Codebase for Large-Scale Vision Models","description":"big_vision is Google Research's official codebase for training large-scale vision models using Jax/Flax. It has been instrumental in developing prominent architectures like Vision Transformer, SigLIP, and MLP-Mixer. This repository offers a robust starting point for researchers to conduct scalable vision experiments on GPUs and Cloud TPUs, scaling seamlessly from single cores to distributed setups.","github":"https://github.com/google-research/big_vision","url":"https://osrepos.com/repo/google-research-big_vision","source":"osrepos.com","sourceDescription":"This repository profile is provided by osrepos.com, an open source repository discovery platform.","repositoryProfile":"https://osrepos.com/repo/google-research-big_vision","generatedFor":"open source discovery and AI-assisted research","markdown":"https://osrepos.com/repo/google-research-big_vision.md","json":"https://osrepos.com/repo/google-research-big_vision.json","topics":["Python","Jupyter Notebook","Machine Learning","Computer Vision","Deep Learning","Vision Transformer","Google Research","AI"],"keywords":["Python","Jupyter Notebook","Machine Learning","Computer Vision","Deep Learning","Vision Transformer","Google Research","AI"],"stars":null,"summary":"big_vision is Google Research's official codebase for training large-scale vision models using Jax/Flax. It has been instrumental in developing prominent architectures like Vision Transformer, SigLIP, and MLP-Mixer. This repository offers a robust starting point for researchers to conduct scalable vision experiments on GPUs and Cloud TPUs, scaling seamlessly from single cores to distributed setups.","content":"## Introduction\n\nbig_vision is the official codebase from Google Research, designed for training large-scale vision models. Built upon the powerful Jax and Flax libraries, it leverages tf.data and TensorFlow Datasets for creating scalable and reproducible input pipelines. This repository serves two primary purposes: to publish the code for various research projects developed within Google, and to provide a robust starting point for researchers to conduct large-scale vision experiments. It supports seamless scaling from a single GPU core to distributed setups with up to 2048 Cloud TPU cores. Notable architectures developed using big_vision include Vision Transformer (ViT), SigLIP, MLP-Mixer, and LiT. While the codebase is dynamic, its core functionality is maintained, and it aims to support Google's internal research, with external contributions generally requiring pre-approval.\n\n## Installation\n\nTo get started with big_vision on a GPU machine, follow these steps. It is highly recommended to use a virtual environment for dependency management.\n\n1.  **Clone the Repository:**\n    bash\n    git clone https://github.com/google-research/big_vision\n    cd big_vision/\n    \n\n2.  **Install Python Dependencies:**\n    bash\n    pip3 install --upgrade pip\n    pip3 install -r big_vision/requirements.txt\n    \n\n3.  **Install JAX with CUDA Support:**\n    bash\n    pip3 install --upgrade \"jax[cuda]\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n    \n    *Note: You may need to adjust the JAX package based on your specific CUDA and cuDNN versions. Refer to the official JAX documentation for more details.*\n\n4.  **Prepare TFDS Data:**\n    big_vision uses `tensorflow_datasets` for unified and reproducible access to standard datasets. It's recommended to prepare datasets separately before running experiments. For example, to download and preprocess `cifar100`, `oxford_iiit_pet`, and `imagenet_v2`:\n    bash\n    cd big_vision/\n    python3 -m big_vision.tools.download_tfds_datasets cifar100 oxford_iiit_pet imagenet_v2\n    \n    Some datasets, like `imagenet2012`, require manual download of raw data files into `$TFDS_DATA_DIR/downloads/manual/` before running the `download_tfds_datasets` command.\n\n## Examples\n\nOnce installed and data is prepared, you can run training jobs. Here are examples for GPU machines:\n\n*   **Train ViT-S/16 on ImageNet:**\n    bash\n    python3 -m big_vision.train --config big_vision/configs/vit_s16_i1k.py --workdir workdirs/`date '+%m-%d_%H%M'`\n    \n\n*   **Train MLP-Mixer-B/16 (with GPU-specific batch size):**\n    bash\n    python3 -m big_vision.train --config big_vision/configs/mlp_mixer_i1k.py:gpu8 --workdir workdirs/`date '+%m-%d_%H%M'`\n    \n    The repository's README also provides detailed instructions and commands for running experiments on Google Cloud TPU VMs, including multi-host setups and FSDP training.\n\n## Why Use big_vision?\n\nbig_vision offers several compelling advantages for researchers and developers working with large-scale vision models:\n\n*   **Scalability:** Designed for high-performance training, it scales effortlessly from single GPU machines to massive distributed setups on Cloud TPUs, supporting up to 2048 TPU cores.\n*   **Research Foundation:** It is the foundational codebase for numerous cutting-edge research projects from Google, including Vision Transformer, MLP-Mixer, LiT, and SigLIP, providing battle-tested implementations.\n*   **Robustness:** Training jobs are robust to interruptions, capable of seamlessly resuming from the last saved checkpoint, ensuring reliability for long-running experiments.\n*   **Powerful Configuration System:** Features a flexible configuration system that allows for easy customization and extension of training parameters and modules.\n*   **Reproducibility:** Utilizes `tf.data` and `TensorFlow Datasets` to ensure scalable and reproducible input pipelines, crucial for scientific research.\n\n## Links\n\n*   **GitHub Repository:** [https://github.com/google-research/big_vision](https://github.com/google-research/big_vision)\n*   **ViT Baseline Paper:** [Better plain ViT baselines for ImageNet-1k](https://arxiv.org/abs/2205.01580)\n*   **JAX Documentation:** [https://github.com/jax-ml/jax#pip-installation-gpu-cuda](https://github.com/jax-ml/jax#pip-installation-gpu-cuda)\n*   **TensorFlow Datasets Catalog:** [https://www.tensorflow.org/datasets/catalog/overview#all_datasets](https://www.tensorflow.org/datasets/catalog/overview#all_datasets)","metrics":{"detailViews":13,"githubClicks":12},"dates":{"published":null,"modified":"2025-12-31T00:00:59.000Z"}}