OneScience-Group/Earthformer

Earthformer

Model Introduction

Earthformer was proposed by researchers from Amazon Web Services (AWS) in collaboration with the Hong Kong University of Science and Technology. It aims to address the prohibitive computational cost of conventional Transformers when processing high-dimensional geophysical data.

Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

https://arxiv.org/abs/2207.05833

Model Description

Earthformer is a space-time Transformer model designed for Earth system forecasting, such as weather and climate. Its core component is a novel attention mechanism called Cuboid Attention.

Use Cases

Scenario	Description
Weather Forecast Training	Train the Earthformer weather forecasting model using SEVIR structured data.
Local Quick Validation	Use synthetic data to verify data loading, model training, inference, and inference result visualization.
ModelScope / OneCode Execution	Download as a standalone model package, install dependencies, and run scripts directly.
Multi-GPU Training	Launch multi-process training via torchrun.

Usage Guide

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

The commands below assume they are run from the Earthformer project root directory.

Hardware Requirements

• Training and inference require a GPU or DCU recognized by PyTorch; CPU can be used to generate synthetic data and verify configuration, but cannot run the current training and inference scripts.
• Multi-GPU training uses the NCCL backend. Please make sure the device driver, communication libraries, and PyTorch version are compatible.
• DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

Download the Model Package

hf download OneScience-Group/Earthformer --local-dir ./Earthformer
cd Earthformer

Install the Runtime Environment

DCU Environment

# Please activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Please activate CONDA first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Generate Synthetic Data

Synthetic data is only used to verify the data protocol and program flow; it does not represent real SEVIR data or the model's forecast quality:

python script/fake_data.py

By default this generates data/synthetic_sevir/{train,val,test}.npz and metadata.json. To generate synthetic data matching the official SEVIR spatial size:

python script/fake_data.py --output-dir data/synthetic_sevir_384 --height 384 --width 384

Training

Single GPU:

python script/train.py

Multi-GPU:

torchrun --nproc_per_node=8 script/train.py

Training starts from random initialization and saves to data/checkpoint/earthformer.pt by default.

Training Weights

This repository provides weights trained on SEVIR data in the weights/ folder. The weight files will be uploaded soon and are expected to be available in the near future.

Inference

Inference reads the training checkpoint by default and writes the output to output/predictions.npz.

python script/inference.py

Evaluation and Visualization

python script/result.py

OneScience Official Information

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation & License

• This repository is the OneScience reproduction of the original Earthformer paper. The official Earthformer implementation is released under the Apache License 2.0; the use of this repository's code and the SEVIR data remains subject to the licenses and terms of use of the respective projects.