Clay Foundation Model

An open source AI model and interface for Earth.

Getting started

Quickstart

Launch into a JupyterLab environment on

Binder	Planetary Computer	SageMaker Studio Lab

Installation

Basic

To help out with development, start by cloning this repo-url

git clone <repo-url>
cd model

Then we recommend using mamba to install the dependencies. A virtual environment will also be created with Python and JupyterLab installed.

mamba env create --file environment.yml

[!NOTE] The command above will only work for Linux devices with CUDA GPUs. For installation on macOS devices (either Intel or ARM chips), follow the 'Advanced' section in https://clay-foundation.github.io/model/installation.html#advanced

Activate the virtual environment first.

mamba activate claymodel

Finally, double-check that the libraries have been installed.

mamba list

Usage

Running jupyter lab

mamba activate claymodel
python -m ipykernel install --user --name claymodel  # to install virtual env properly
jupyter kernelspec list --json                       # see if kernel is installed
jupyter lab &

Running the model

The neural network model can be ran via LightningCLI v2. To check out the different options available, and look at the hyperparameter configurations, run:

python trainer.py --help
python trainer.py test --print_config

To quickly test the model on one batch in the validation set:

python trainer.py validate --trainer.fast_dev_run=True

To train the model for a hundred epochs:

python trainer.py fit --trainer.max_epochs=100

To generate embeddings from the pretrained model's encoder on 1024 images (stored as a GeoParquet file with spatiotemporal metadata):

python trainer.py predict --ckpt_path=checkpoints/last.ckpt \
                          --data.batch_size=1024 \
                          --data.data_dir=s3://clay-tiles-02 \
                          --trainer.limit_predict_batches=1

More options can be found using python trainer.py fit --help, or at the LightningCLI docs.