Earthformer

By Zhihan Gao, Xingjian Shi, Hao Wang, Yi Zhu, Yuyang Wang, Mu Li, Dit-Yan Yeung.

This repo is the official implementation of "Earthformer: Exploring Space-Time Transformers for Earth System Forecasting" that will appear in NeurIPS 2022.

Check our poster.

Tutorials

• Inference Tutorial of Earthformer on EarthNet2021.

Introduction

Conventionally, Earth system (e.g., weather and climate) forecasting relies on numerical simulation with complex physical models and are hence both expensive in computation and demanding on domain expertise. With the explosive growth of the spatiotemporal Earth observation data in the past decade, data-driven models that apply Deep Learning (DL) are demonstrating impressive potential for various Earth system forecasting tasks. The Transformer as an emerging DL architecture, despite its broad success in other domains, has limited adoption in this area. In this paper, we propose Earthformer, a space-time Transformer for Earth system forecasting. Earthformer is based on a generic, flexible and efficient space-time attention block, named Cuboid Attention. The idea is to decompose the data into cuboids and apply cuboid-level self-attention in parallel. These cuboids are further connected with a collection of global vectors.

Earthformer achieves strong results in synthetic datasets like MovingMNIST and N-body MNIST dataset, and also outperforms non-Transformer models (like ConvLSTM, CNN-U-Net) in SEVIR (precipitation nowcasting) and ICAR-ENSO2021 (El Nino/Southern Oscillation forecasting).

Cuboid Attention Illustration

Installation

We recommend managing the environment through Anaconda.

First, find out where CUDA is installed on your machine. It is usually under /usr/local/cuda or /opt/cuda.

Next, check which version of CUDA you have installed on your machine:

nvcc --version

Then, create a new conda environment:

conda create -n earthformer python=3.9
conda activate earthformer

Lastly, install dependencies. For example, if you have CUDA 11.6 installed under /usr/local/cuda, run:

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
cd ROOT_DIR/earth-forecasting-transformer
python3 -m pip install -U -e . --no-build-isolation

# Install Apex
CUDA_HOME=/usr/local/cuda python3 -m pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extension git+https://github.com/NVIDIA/apex.git

If you have CUDA 11.7 installed under /opt/cuda, run:

python3 -m pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 -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
cd ROOT_DIR/earth-forecasting-transformer
python3 -m pip install -U -e . --no-build-isolation

# Install Apex
CUDA_HOME=/opt/cuda python3 -m pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extension git+https://github.com/NVIDIA/apex.git

Dataset

MovingMNIST

We follow Unsupervised Learning of Video Representations using LSTMs (ICML2015) to use MovingMNIST that contains 10,000 sequences each of length 20 showing 2 digits moving in a $64\times 64$ frame.

Our MovingMNIST DataModule automatically downloads it to datasets/moving_mnist.

N-body MNIST

The underlying dynamics in the N-body MNIST dataset is governed by the Newton's law of universal gravitation:

$\frac{d^2\boldsymbol{x}_{i}}{dt^2} = - \sum_{j\neq i}\frac{G m_j (\boldsymbol{x}_{i}-\boldsymbol{x}_{j})}{(|\boldsymbol{x}_i-\boldsymbol{x}_j|+d_{\text{soft}})^r}$

where $\boldsymbol{x}_{i}$ is the spatial coordinates of the $i$-th digit, $G$ is the gravitational constant, $m_j$ is the mass of the $j$-th digit, $r$ is a constant representing the power scale in the gravitational law, $d_{\text{soft}}$ is a small softening distance that ensures numerical stability.

The N-body MNIST dataset we used in the paper can be downloaded from https://earthformer.s3.amazonaws.com/nbody/nbody_paper.zip .

In addition, you can also use the following script for downloading / extracting the data:

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

Alternatively, run the following commands to generate N-body MNIST dataset.

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

SEVIR

Storm EVent ImageRy (SEVIR) dataset is a spatiotemporally aligned dataset containing over 10,000 weather events. We adopt NEXRAD Vertically Integrated Liquid (VIL) mosaics in SEVIR for benchmarking precipitation nowcasting, i.e., to predict the future VIL up to 60 minutes given 65 minutes context VIL. The resolution is thus $13\times 384\times 384\rightarrow 12\times 384\times 384$.

To download SEVIR dataset from AWS S3, run:

cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/sevir/download_sevir.py --dataset sevir

A visualization example of SEVIR VIL sequence:

ICAR-ENSO

ICAR-ENSO consists of historical climate observation and stimulation data provided by Institute for Climate and Application Research (ICAR). We forecast the SST anomalies up to 14 steps (2 steps more than one year for calculating three-month-moving-average), given a context of 12 steps (one year) of SST anomalies observations.

To download the dataset, you need to follow the instructions on the official website. You can download a zip-file named enso_round1_train_20210201.zip. Put it under ./datasets/ and extract the zip file with the following command:

unzip datasets/enso_round1_train_20210201.zip -d datasets/icar_enso_2021

EarthNet2021

You may follow the official instructions for downloading EarthNet2021 dataset. We recommend download it via the earthnet_toolket.

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

Alternatively, you may download EarthNet2021x dataset, which is the same as EarthNet2021 dataset except for the file format (.npz for EarthNet2021 and .nc for EarthNet2021x).

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

It requires 455G disk space in total.

Earthformer Training

Find detailed instructions in the corresponding training script folder

• N-body MNIST
• SEVIR&SEVIR-LR
• ENSO
• EarthNet2021

Training Script and Pretrained Models

Find detailed instructions in how to train the models or running inference with our pretrained models in the corresponding script folder.

Dataset	Script Folder	Pretrained Weights	Config
SEVIR	scripts	link	config
ICAR-ENSO	scripts	link	config
EarthNet2021	scripts	link	config
N-body MNIST	scripts	-	-

Citing Earthformer

@inproceedings{gao2022earthformer,
  title={Earthformer: Exploring Space-Time Transformers for Earth System Forecasting},
  author={Gao, Zhihan and Shi, Xingjian and Wang, Hao and Zhu, Yi and Wang, Yuyang and Li, Mu and Yeung, Dit-Yan},
  booktitle={NeurIPS},
  year={2022}
}

Security

See CONTRIBUTING for more information.

Credits

Third-party libraries:

• PyTorch
• PyTorch Lightning
• Apex
• OpenCV
• TensorBoard
• OmegaConf
• YACS
• Pillow
• scikit-learn

License

This project is licensed under the Apache-2.0 License.