DOCSEARCHbyAlgolia
Submit Docs

Clay Foundation Model

repository·main·Indexed 20 days ago

https://github.com/clay-foundation/model

An open-source AI model and interface for Earth observation tasks, providing a framework for training and running foundation models using PyTorch Lightning. The library includes the ClayMAEModule for Vision Transformer-based architectures with various scales (tiny, small, base, large), the ClayDataModule for handling satellite imagery, and specialized tools like EODataset for .npz files and ClaySampler for platform-balanced batching.

Tokens
38.2K
Snippets
122
Records
143
Agent score
60%

What's inside claymodel

  1. How to use the Clay Foundation Model

    main

    The Clay model is a Vision Transformer adapted for Earth Observation (EO) data. It takes satellite imagery, location, and time as inputs to produce mathematical embeddings. There are three primary usage patterns:

    1. Generate semantic embeddings: Use embeddings to find specific features (e.g., surface mines, aquaculture, or concentrated animal feeding operations) at any given location and time.
    2. Fine-tune for downstream tasks: Leverage pre-trained knowledge to perform specific tasks more efficiently than training from scratch. This includes:
      • Classification/Regression: Identifying crop types, land cover, or predicting variables like above-ground biomass.
      • Change Detection: Detecting deforestation, wildfires, flooding, or urban development by training a model to process embeddings over time.
    3. Use as a backbone: Utilize the model as the foundational architecture for other specialized models.

  2. Understand the Clay v0 Model Overview

    main

    Clay v0 is a self-supervised modified vision transformer (ViT) model trained as a Masked Autoencoder (MAE). It is designed to process stacks of geospatial data and output high-dimensional vector embeddings that capture spatial, temporal, and spectral information about Earth.

    Key Capabilities:

    • Input: Stacks of geospatial data (Sentinel-2, Sentinel-1, and DEM).
    • Output: Vector embeddings representing specific areas of Earth at specific points in time.
    • Architecture: A modified ViT encoder for generating embeddings and a decoder for reconstructing masked parts of the original image.

    Model Specifications:

  3. Fine-tune a segmentation head on top of Clay

    main

    You can fine-tune a segmentation head using feature maps extracted from a frozen Clay encoder. This process uses the Segmentor class to extract feature maps and upsample them to the original image size using convolution and pixel shuffle operations. The implementation typically uses PyTorch Lightning for the training workflow.

    # The Segmentor class is used to extract feature maps from a frozen Clay encoder
# and upsample them to the original image size.

  4. Understand the finetuning architecture vs training from scratch

    main

    When performing finetuning with Clay, the architecture differs significantly from training a model from scratch:

    Finetuning Architecture

    In finetuning, the encoder is absent from the downstream model definition. Instead, the pre-trained Clay model's encoder is used to generate embeddings, which are then passed to a decoder network.

    1. Embedding Reshaping: Embeddings are reshaped from batch size * (band groups length * number of patches) * embedding size to batch size * (band groups length * embedding size) * patch height * patch width.
    2. Decoder: The decoder typically consists of a series of ConvTranspose2d, Upsample, and ReLU layers to map embeddings back to the spatial dimensions of the labels.

    Training from Scratch Architecture

    When training from scratch, the model includes a dedicated encoder (e.g., Conv2d and MaxPool2d layers) to learn latent information from the raw pixels, followed by a decoder.

    Common Components

    For binary segmentation tasks, both architectures use:

    • Loss Function: binary_cross_entropy_with_logits.
    • Output Processing: sigmoid and max functions applied to predictions to obtain final segmentation results.
    • Evaluation Metrics: Dice coefficient, Intersection over Union (IoU), F1 score, precision, and recall.

  5. Configure Input Data Modalities for Clay v0

    main

    To use the Clay v0 model, your input data must follow a fixed specification of 13 bands. Each data entry is a stack containing:

    • Sentinel-2: 10 bands
    • Sentinel-1: 2 bands
    • DEM (Digital Elevation Model): 1 band

    The model was trained using 3 timesteps of data for each location. The training data used 10km x 10km MGRS tiles, which were split into 512 x 512 chips.

  6. Understand the Clay benchmark dataset (Cloud to Street - Microsoft flood dataset)

    main

    Clay evaluates its foundation model using benchmark datasets with suitable downstream tasks. The primary benchmark used for initial linear probing and fine-tuning experiments is the Cloud to Street - Microsoft flood dataset.

    Task and Data Structure

    • Task: Semantic segmentation of water pixels associated with recorded flood events.
    • Input Format: Datacubes consisting of three inputs at 512x512 pixel resolution:
      1. Sentinel-2: Surface reflectance (L2A / Bottom-of-Atmosphere) queried via Microsoft Planetary Computer STAC API.
      2. Sentinel-1: VV and VH polarizations.
      3. Copernicus Digital Elevation Model (DEM).
    • Metadata: Datacubes preserve geospatial coordinates, timestamps, and flood event information to ensure compatibility with the Clay Foundation model's embedding requirements.

    Data Generation Requirements

    To generate these datacubes using the provided pipeline, you require:

    • An AWS account (for writing datacubes to S3).
    • A Microsoft Planetary Computer API Key (for querying STAC APIs).

  7. Configure embedding granularity with embeddings_level

    main

    The --model.embeddings_level flag determines the shape and granularity of the generated embeddings:

    • mean (default): Calculates the average across the patch dimension. Results in one embedding per MGRS tile with a size of 768.
    • patch: Keeps embeddings at the patch level. The embedding array size is 16 * 16 * 768, representing one embedding per patch.
    • group: Keeps the full dimensionality of the encoder output, including the band group dimension. The array size is 6 * 16 * 16 * 768.

  8. Perform linear probing to analyze Clay representations

    main

    Linear probing is a technique used to explore the representations learned by the Clay Foundation model during its pre-training.

    To implement this, you use a finetuned architecture (a decoder without an encoder) and execute it periodically (e.g., every n epochs) during the Foundation model's training cycle. This is typically implemented using a PyTorch Lightning callback.

  9. How Clay's sensor-agnostic architecture works

    main

    Clay v1.5 is designed to be sensor-agnostic. It can process data from any satellite instrument provided you supply the correct metadata. The model relies on three pieces of information to interpret the input:

    1. Band Order: The sequence of spectral bands in the input tensor.
    2. Wavelengths: The specific central wavelengths for each band (converted to nm).
    3. Normalization Statistics: The mean and standard deviation for each band used to normalize the raw pixel values.

    Metadata for supported sensors (like Sentinel-2, Landsat, NAIP, etc.) is managed in configs/metadata.yaml.

  10. Migrate imports to the claymodel package

    main

    The Clay Foundation Model has moved from a source-based import structure (src.*) to a proper Python package named claymodel. You must update your import statements to use the new package prefix.

    # Old development imports
from src.datamodule import ClayDataModule
from src.module import ClayMAEModule
from src.model import ClayMAEEncoder

# New package imports
from claymodel.datamodule import ClayDataModule
from claymodel.module import ClayMAEModule
from claymodel.model import ClayMAEEncoder

  11. Prepare a Docker image in ECR for the datacube pipeline

    main

    To run the datacube pipeline as an AWS Batch job using a fetch-and-run approach, you must first build and push a Docker image to Amazon Elastic Container Registry (ECR).

    1. Navigate to the batch scripts directory.
    2. Build the image with the tag format: <ecr_repo_id>.dkr.ecr.<region>.amazonaws.com/fetch-and-run.
    3. Authenticate Docker with ECR.
    4. Push the image.

    Note: Replace 12345 with your actual ecr_repo_id and ensure your AWS profile (clay) is configured correctly.

    ecr_repo_id=12345
cd scripts/pipeline/batch
docker build -t $ecr_repo_id.dkr.ecr.us-east-1.amazonaws.com/fetch-and-run .

aws ecr get-login-password --profile clay --region us-east-1 | docker login --username AWS --password-stdin $ecr_repo_id.dkr.ecr.us-east-1.amazonaws.com

docker push $ecr_repo_id.dkr.ecr.us-east-1.amazonaws.com/fetch-and-run

  12. Prepare the directory structure for custom AOI processing

    main

    Before running the model over custom Areas of Interest (AOIs), you must create a specific directory structure within the model repository to hold MGRS tiles, imagery chips, and embeddings.

    Run the following commands to initialize the required folders:

    # Move into the model repository
cd /path/to/repository/model/

# Ensure data sub-directories exist
mkdir data/mgrs
mkdir data/chips
mkdir data/embeddings