DOCSEARCHbyAlgolia
Submit Docs

Earthformer: Earth Forecasting Transformer

repository·main·Indexed 19 days ago

https://github.com/amazon-science/earth-forecasting-transformer

A space-time Transformer architecture designed for Earth system forecasting, such as weather and climate. It features a 'Cuboid Attention' mechanism to efficiently process spatiotemporal data. The repository includes training scripts and support for datasets including SEVIR, ICAR-ENSO, EarthNet2021, N-body MNIST, and Moving MNIST.

Tokens
15.1K
Snippets
52
Records
59
Agent score
66%

What's inside Earthformer

  2. Install Earthformer via Conda

    main

    Earthformer is recommended to be managed via Anaconda. The installation process involves creating a Python 3.9 environment and installing specific versions of PyTorch and PyTorch Lightning compatible with your CUDA version.

    Note: You must identify your CUDA installation path (e.g., /usr/local/cuda or /opt/cuda) and check your version using nvcc --version before proceeding.

    # 1. Create and activate environment
conda create -n earthformer python=3.9
conda activate earthformer

# 2. Install dependencies (Example for CUDA 11.6)
python3 -m pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html
python3 -m pip install pytorch_lightning==1.6.4
python3 -m pip install xarray netcdf4 opencv-python earthnet==0.3.9

# 3. Install Earthformer in editable mode
cd ROOT_DIR/earth-forecasting-transformer
python3 -m pip install -U -e . --no-build-isolation

  3. Train Earthformer on the ICAR-ENSO dataset

    main

    To train Earthformer using the ICAR-ENSO dataset, execute the train_cuboid_enso.py script. You must configure the training parameters in cfg.yaml before running. The command requires setting MASTER_ADDR and MASTER_PORT environment variables for distributed training coordination.

    MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_enso.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_enso

  4. Download and prepare N-body MNIST dataset

    main

    You can prepare the N-body MNIST dataset using either a download script or by generating it locally.

    Option 1: Download existing data

    cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/nbody/download_nbody_paper.py

    Option 2: Generate dataset from scratch

    cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/nbody/generate_nbody_dataset.py --cfg ./scripts/datasets/nbody/cfg.yaml
    # Download existing data
python ./scripts/datasets/nbody/download_nbody_paper.py

# Or generate from scratch
python ./scripts/datasets/nbody/generate_nbody_dataset.py --cfg ./scripts/datasets/nbody/cfg.yaml

  7. Train Earthformer on EarthNet2021 with auxiliary meso scale data

    main

    To train Earthformer on the EarthNet2021 dataset using auxiliary meso scale data, execute the train_cuboid_earthnet.py script. You must configure the training parameters in the cfg.yaml file before running the command.

    Required environment variables for distributed training:

    • MASTER_ADDR: The address of the master node (e.g., localhost).
    • MASTER_PORT: The port for communication (e.g., 10001).

    Arguments:

    • --gpus: Number of GPUs to use.
    • --cfg: Path to the configuration YAML file.
    • --ckpt_name: The name of the checkpoint to use.
    • --save: The directory where results will be saved.
    MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_earthnet.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_earthnet_w_meso

  9. Train Earthformer on the Moving MNIST dataset

    main

    To train the Earthformer model using the Moving MNIST dataset, execute the train_cuboid_mnist.py script. You must provide the number of GPUs, a configuration file, a checkpoint name, and a save directory.

    Before running, ensure you have updated the settings in cfg.yaml to match your requirements. The training command requires setting MASTER_ADDR and MASTER_PORT environment variables for distributed training coordination.

    MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_mnist.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_mnist

  10. Generate the N-body MNIST dataset

    main

    If you prefer to generate the N-body MNIST dataset locally instead of downloading it, use the generate_nbody_dataset.py script. You must provide a configuration file using the --cfg flag to define the generation parameters.

    cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/nbody/generate_nbody_dataset.py --cfg ./scripts/datasets/nbody/cfg.yaml

  11. Download and prepare EarthNet2021 dataset

    main

    It is recommended to download EarthNet2021 using the earthnet_toolkit. You can download either the standard earthnet2021 dataset or the earthnet2021x dataset (which uses .nc format instead of .npz).

    Note: This dataset requires approximately 455GB of disk space.

    Using earthnet_toolkit:

    import earthnet as en
# For standard EarthNet2021
en.download(dataset="earthnet2021", splits="all", save_directory="./datasets/earthnet2021")

# For EarthNet2021x
en.download(dataset="earthnet2021x", splits="all", save_directory="./datasets/earthnet2021x")
    import earthnet as en
en.download(dataset="earthnet2021", splits="all", save_directory="./datasets/earthnet2021")