README.md

---
datasets:
- SEVIR
frameworks:
- PyTorch
language:
- en
- zh
license: apache-2.0
tags:
- OneScience
- Earth Science
- Weather Forecast
- Spatiotemporal Forecast
- SEVIR
- Earthformer
tasks: []
---

<p align="center">
  <strong>
    <span style="font-size: 30px;">Earthformer</span>
  </strong>
</p>

# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 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

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

### Install the Runtime Environment

**DCU Environment**

```bash
# 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**

```bash
# 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:

```bash
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:

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

### Training

Single GPU:

```bash
python script/train.py
```

Multi-GPU:

```bash
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`.

```bash
python script/inference.py
```

### Evaluation and Visualization

```bash
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.