RADIO (Agglomerative Vision Foundation Model)

repository·main·Indexed 23 days ago

https://github.com/nvlabs/radio

A PyTorch implementation of vision foundation models designed to reduce multiple domains into a single representation. It features models like C-RADIOv4 and supports multimodal alignment via teacher adaptors (e.g., siglip2-g, clip, sam). The library provides tools for zero-shot ImageNet classification, K-NN classification, patch position prediction, and semantic segmentation linear probing on ADE20k. Models can be loaded via TorchHub or HuggingFace transformers.

Tokens
4.7K
Snippets
14
Records
21
Agent score
83%

What's inside RADIO

  1. Understand RADIO1D (Research) architecture

    main

    RADIO1D is a Vision Transformer variant that compresses spatial tokens into a variable-length 1D sequence of "global tokens" during encoding, which can then be reconstructed via a decoder.

    It exposes two "necks":

    • encoder: Compressed 1D global tokens, shape (B, num_tokens, C).
    • decoder: Spatially-reconstructed features, shape (B, H*W, C).

    You can control the number of tokens and which neck to return using these arguments in radio_model() or RADIOModel.forward():

    • num_tokens (Optional[int]): The number of tokens to keep in the 1D encoder output (defaults to the model's maximum).
    • neck_name (Optional[str]): Which neck's output to return (defaults to returning a dictionary of all necks for multi-neck models).
  2. How RADIOv2.5 handles input resolution and mode switching

    main

    Unlike RADIOv2, which suffered from "mode switching" (where features changed drastically depending on whether the input resolution was below or above ~700px), RADIOv2.5 models are designed to be resolution-agnostic.

    Key improvements in v2.5:

    • Consistent Representations: Features remain stable across different resolutions, allowing for consistent behavior in tasks like zero-shot classification and VLLM integration.
    • High-Resolution Performance: RADIOv2.5 performs exceptionally well at high resolutions (>= 768px).
    • Simultaneous Feature Extraction: Because the model no longer switches modes, you can simultaneously extract CLIP, DINOv2, and SAM features from the same high-resolution image, and all will be meaningful.
  3. Load RADIO models via TorchHub

    main

    You can load RADIO models using torch.hub.load. For the latest C-RADIOv4 models, use model_version='c-radio_v4-h' or model_version='c-radio_v4-so400m'.

    Important: If you have used TorchHub with this repository in the past, it is strongly recommended to run with force_reload=True once to ensure you pull the latest source code.

    import torch
    
    model_version = "c-radio_v4-h"
    model = torch.hub.load('NVlabs/RADIO', 'radio_model', version=model_version, progress=True, skip_validation=True, force_reload=True)
    model.cuda().eval()
  4. Perform K-NN Classification

    main

    K-NN classification involves embedding summary vectors from the model into a database and classifying new vectors by computing a weighted sum of the $k$ nearest neighbors using cosine similarity. This follows the weighting algorithm used in DINO and DINOv2 (with $\tau = 0.07$).

    Use the knn_classification.py script, which integrates with the HuggingFace Datasets library. This script supports single-GPU and multi-GPU (via torchrun) configurations.

    Resolution Options:

    • Default: Resizes shorter dimension to 378px and center-crops the larger dimension to 378px.
    • Custom Square: Use --resolution 224 224 (or any value divisible by the model's patch size).
    • Non-square: Use a single value like --resolution 378 to resize the smaller dimension and preserve aspect ratio.
    # Food101
    python knn_classification.py --dataset food101
    
    # Oxford Pets
    python knn_classification.py --dataset "jonathancui/oxford-pets" --eval-split test
    
    # ImageNet-1K (Multi-GPU)
    torchrun --nproc-per-node 8 knn_classification.py --dataset imagenet-1k
    
    # ImageNet Sketch (using ImageNet-1K training database)
    torchrun --nproc-per-node 8 knn_classification.py --dataset imagenet-1k --eval-dataset imagenet_sketch --eval-split train
  5. Load RADIO models via HuggingFace

    main

    You can load RADIO models from HuggingFace using the transformers library. Note that for C-RADIO models, you should use the appropriate repository name (e.g., nvidia/C-RADIOv4-H). For E-RADIO or RADIO models, use their respective repository names.

    To use adaptors (like clip or sam), you must first load the configuration, set the adaptor_names list, and then pass that config to AutoModel.from_pretrained.

    import torch
    from PIL import Image
    from transformers import AutoModel, CLIPImageProcessor, AutoConfig
    
    hf_repo = "nvidia/C-RADIOv4-H"
    
    # Basic loading
    image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
    model = AutoModel.from_pretrained(hf_repo, trust_remote_code=True)
    model.eval().cuda()
    
    # Loading with adaptors
    config = AutoConfig.from_pretrained(hf_repo, trust_remote_code=True)
    config.adaptor_names = ["clip", "sam"]
    model = AutoModel.from_pretrained(hf_repo, trust_remote_code=True, config=config)
    model.eval().cuda()
    
    # Inference
    image = Image.open('./assets/radio.png').convert('RGB')
    pixel_values = image_processor(images=image, return_tensors='pt', do_resize=True).pixel_values.cuda()
    
    # Accessing adaptor-specific outputs
    clip_summary, clip_features = model(pixel_values)["clip"].summary, model(pixel_values)["clip"].features
    sam_summary, sam_features = model(pixel_values)["sam"].summary, model(pixel_values)["sam"].features
  6. Setup RADIO Semantic Segmentation Linear Probing

    main

    To perform semantic segmentation on the ADE20k dataset using linear probing on frozen RADIO features, follow these setup steps:

    1. Install Dependencies: Install the required Python packages using the provided requirements file.
    2. Prepare Dataset: Download the ADE20k dataset following the official instructions at https://groups.csail.mit.edu/vision/datasets/ADE20K/.
    3. Authenticate Hugging Face: Ensure you are logged into the Hugging Face Hub to access necessary models or assets.

    This process was verified on an 8xA100 system using the nvcr.io/nvidia/pytorch:23.11-py3 container.

    pip install -r requirements.txt
    
    huggingface-cli login
  7. Load RADIOv2.5 models via TorchHub

    main

    You can load the RADIOv2.5 model family using torch.hub.load. The API is consistent with previous versions.

    Available versions:

    • radio_v2.5-h: The best performing model (ViT-H/16).
    • radio_v2.5-l: A high-quality, faster, and smaller model (ViT-L/16). Recommended as a replacement for RADIOv2.
    • radio_v2.5-b: A smaller and faster model (ViT-B/16).

    Note: When requesting these models for the first time, set force_reload=True to ensure you pull the latest source code and weights.

  8. Perform Zero-Shot ImageNet Classification

    main

    Since RADIO learns to match CLIP teacher features, you can replace the CLIP vision tower with RADIO while using the pre-trained CLIP text tower for zero-shot classification.

    Use the zero_shot_imagenet.py script to evaluate. You can run on a single GPU or use torchrun for multi-GPU setups. You can also specify different datasets (like imagenet_sketch) and resolutions.

    Resolution Behavior:

    • Providing a single value (e.g., --resolution 378) resizes the shorter dimension to that value while preserving aspect ratio (non-square).
    • Providing two values (e.g., --resolution 378 378) results in square crops.
    • Note: Either or both dimensions must be divisible by 14 for the specified resolution to work (up to 1050px).
    # Single GPU
    python zero_shot_imagenet.py
    
    # Multi-GPU
    torchrun --nproc-per-node 8 zero_shot_imagenet.py
    
    # Classify ImageNet-Sketch with specific resolution
    torchrun --nproc-per-node 8 zero_shot_imagenet.py --resolution 378 --dataset imagenet_sketch --split train
    
    # Square crops
    torchrun --nproc-per-node 8 zero_shot_imagenet.py --resolution 378 378
  9. Train a linear head for semantic segmentation

    main

    To start the training process for linear probing, you must first set the ADE20K_ROOT_DIR environment variable to the path where your ADE20k dataset is located.

    Then, run the training script using torch.distributed.launch. You must pass the dataset root directory to both the train_dataloader and val_dataloader via the --cfg-options flag.

    export ADE20K_ROOT_DIR=/path/to/ade20k/dataset
    
    python -m torch.distributed.launch --nnodes=1 --nproc_per_node=8 train.py configs/radio/radio_linear_8xb2-80k_ade20k-512x512.py --launcher pytorch --cfg-options "train_dataloader.dataset.data_root=/${ADE20K_ROOT_DIR}" --cfg-options "val_dataloader.dataset.data_root=${ADE20K_ROOT_DIR}"
  10. Use the SigLIP adaptor in RADIO

    main

    RADIO supports using the SigLIP adaptor (based on ViT-SO400M-14-SigLIP-384 from OpenCLIP) instead of the default DFN CLIP.

    To use it, pass the --adaptor-name siglip argument to the relevant script (e.g., examples/zero_shot_imagenet.py). Note that classification results with the SigLIP head may be slightly lower than DFN CLIP, so use it primarily for compatibility requirements.

    --adaptor-name siglip
  11. Run RADIO with mixed precision (AMP)

    main

    For improved performance, especially when using self-attention, you can run RADIO using PyTorch's Automatic Mixed Precision (AMP).

    with torch.autocast('cuda', dtype=torch.bfloat16):
        summary, spatial_features = model(x)