DOCSEARCHbyAlgolia
Submit Docs

big_vision

repository·main·Indexed 25 days ago

https://github.com/google-research/big_vision

A high-performance research codebase for training large-scale vision and multimodal models using JAX, Flax, and TensorFlow Datasets, optimized for Cloud TPUs and GPUs. It includes implementations and configurations for models such as CapPa, CLIPPO, BiT, FlexiViT, and GIVT.

Tokens
21.1K
Snippets
48
Records
106
Agent score
86%

What's inside big_vision

  1. Overview of JetFormer

    main
    JetFormer is an autoregressive decoder-only transformer designed for joint generative modeling of raw images and text. Unlike models that rely on separately pretrained modality-specific encoders and decoders, JetFormer is trained to directly maximize the likelihood of raw data. It utilizes a normalizing flow model to obtain soft-token image representations, which serves as both an image encoder for perception and an image decoder for generation. This allows the model to perform high-fidelity text-to-image generation and robust image understanding without relying on complex mixtures of perceptual losses or pretrained autoencoders.

  2. Overview of LiT (Locked-image text Tuning)

    main

    LiT (Zero-Shot Transfer with Locked-image text Tuning) is a method for multimodal research. It allows for zero-shot transfer by tuning text models to vision models while keeping the vision models locked.

    Key resources:

  3. Overview of Big Vision

    main

    Big Vision is a codebase designed for training large-scale vision models on Cloud TPU VMs or GPU machines. It is built using the Jax and Flax libraries and utilizes tf.data and TensorFlow Datasets for scalable, reproducible input pipelines.

    Key features include:

    • Seamless scaling from a single TPU core to distributed setups with up to 2048 TPU cores.
    • Implementation of various research projects in architecture, multimodal learning, and training techniques.
    • Support for large-scale vision experiments out-of-the-box.

  4. Overview of PaliGemma Vision-Language Model

    main

    PaliGemma is an open vision-language model (VLM) inspired by PaLI-3. It is designed for a wide range of vision-language tasks, including:

    • Image and short video captioning
    • Visual Question Answering (VQA)
    • Text reading
    • Object detection
    • Object segmentation

    Architecture Details:

    • Visual Encoder: SigLIP (specifically SigLIP-So400m/14).
    • Language Model: Gemma (PaliGemma uses Gemma 2B; PaliGemma 2 uses Gemma 2 {2B, 9B, 27B}).
    • Mechanism: Images are converted into "soft tokens" by the SigLIP encoder. Input text (the "prefix") is tokenized by Gemma. Image tokens and prefix tokens are concatenated and passed to the Gemma decoder with full block-attention to generate an output text (the "suffix") auto-regressively with masked attention.

  5. Research Projects in Big Vision

    main

    The codebase contains implementations and configurations for numerous research papers across several domains:

    Architecture Research

    Includes implementations for Transformers for Image Recognition (ViT), MLP-Mixer, UViM, FlexiViT, Dual PatchNorm, GIVT, Jet, and JetFormer.

    Multimodal Research

    Includes implementations for LiT (Locked-image Text Tuning), CLIPPO, SigLIP, SigLIP 2, CAPPA, Three Towers, PaLI, PaLI-3, LocCa, and PaliGemma/PaliGemma 2.

    Training Techniques

    Includes implementations for Knowledge Distillation, Sharpness-Aware Minimization (SAM), Surrogate Gap Minimization (GSAM), and VeLO.

  6. Explore Image/text multimodal learning projects

    main

    The big_vision/configs/proj/image_text/ directory contains configurations and Colab notebooks for various image/text multimodal learning projects.

    Key projects available in this directory include:

    • LiT (Locked-image text Tuning): Focuses on zero-shot transfer. See README_lit.md for details.
    • SigLIP 2: Focuses on multilingual vision-language encoders with improved semantic understanding, localization, and dense features. See README_siglip2.md for details.

  7. Prepare TFDS datasets

    main
    Big Vision uses tensorflow_datasets (tfds). You can automatically download and prepare most datasets using the big_vision.tools.download_tfds_datasets module. For datasets requiring manual downloads (like imagenet2012), you must manually place the official files in $TFDS_DATA_DIR/downloads/manual/ before running the preparation tool.

  9. Download PaliGemma checkpoints from Kaggle

    main

    PaliGemma checkpoints are hosted on Kaggle and require a Kaggle account and acknowledgment of the license terms. You must set your KAGGLE_USERNAME and KAGGLE_KEY environment variables to download them via curl.

    Available checkpoint types include:

    • Pretrained (pt): pt-224, pt-448, pt-896. Use these for fine-tuning initialization.
    • Mixture (mix): mix-224, mix-448. These are trained on a mixture of transfer tasks and are designed for out-of-the-box prediction with natural language prompts.
    • Transfers: Specialized checkpoints for specific tasks (e.g., Captioning, QA, Segmentation) available on Kaggle.
    export KAGGLE_USERNAME=
export KAGGLE_KEY=

# Example for downloading a specific model
export MODEL_NAME=paligemma2-3b-pt-224

mkdir ckpts/
cd ckpts/

curl -L -u $KAGGLE_USERNAME:$KAGGLE_KEY\ 
  -o pt_3b_224.bf16.npz \ 
  https://www.kaggle.com/api/v1/models/google/paligemma-2/jax/$MODEL_NAME/1/download/$MODEL_NAME.b16.npz

  10. Train a CLIPPO model on TPU

    main

    To launch CLIPPO training on Google Cloud TPU VMs, use the run_tpu.sh script.

    Note on Data: By default, the configuration trains on MS-COCO Captions (via TFDS) and initializes with ImageNet21k weights for sanity checking. To train on large datasets like LAION-400M or YFCC-100M, you must first wrap your dataset using TensorFlow Datasets (TFDS) and update the config.

    To train with custom data and ImageNet1k evaluations, append the following flags to the config path: --config big_vision/configs/proj/clippo/train_clippo.py:test_with_coco=False,i1k_eval=True

    gcloud alpha 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.contrastive --config big_vision/configs/proj/clippo/train_clippo.py --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'`"

  12. Train GIVT generative models

    main

    To train GIVT (Generative Infinite-Vocabulary Transformers) models on ImageNet 2012, use the big_vision.trainers.proj.givt.generative module. You must provide a config path and a Google Cloud Storage work directory.

    Note: The default givt_imagenet2012.py config uses Imagenette to allow for immediate testing without a manual ImageNet download. This will overfit quickly; for paper-reproducible results, download ImageNet and update your configuration.

    You can override configuration parameters by appending :key=value to the config path.

    python -m big_vision.trainers.proj.givt.generative \
  --config big_vision/configs/proj/givt/givt_imagenet2012.py \
  --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`

# Example with config override:
python -m big_vision.trainers.proj.givt.generative \
  --config big_vision/configs/proj/givt/givt_imagenet2012.py:model_size=large \
  --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`