Skip to content

Repository files navigation

[ICLR2026] Task-Adaptive Parameter-Efficient Fine-Tuning for Weather Foundation Models

Framework: PyTorch

Contents

  1. Installation Instructions
  2. Dataset Preparation
  3. Execution Instructions
  4. Acknowledgement

Installation Instructions

  • We use Python 3.10, PyTorch 2.5.1 (CUDA 11.8 build).

#Make sure that the WeatherPEFT is placed in the current user directory for dataset register.
cd WeatherPEFT

conda create -n WeatherPEFT python=3.10
conda activate WeatherPEFT

pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118
conda install -c conda-forge xarray==2025.1.2 numpy==2.1.3 pandas==2.2.3 dask netCDF4 bottleneck

pip install -r requirements.txt

Dataset Preparation

Download the constant data, including the static variable for Aurora (Land-sea mask, Soil type, Surface-level geopotentia), precomputed statistics for normalization (mean and std), and ERA5-CH precipitation ground-truth and Climatology data for precipitation evaluation. Put them in WeatherPEFT/aux_data

  1. Downscale. In this experiment, we downscale 5.625° ERA5 data to 1.40625° ERA5 data both at a global scale and 6-hour intervals. The pre-processed 5.625° and 1.40625° ERA5 data data can be downloaded in WeatherBench. Please follow dataset structure below:s
    |--WeatherPEFT
        |--datasets
            |--5.625
                |--2m_temperature
                |--10m_u_component_of_wind
                |--10m_v_component_of_wind 
                |--temperature
                |--u_component_of_wind
                |--v_component_of_wind
                |--specific_humidity
                |--geopotential
            |--1.40625
                |--2m_temperature
                |--10m_u_component_of_wind
                |--10m_v_component_of_wind 
                |--temperature
                |--u_component_of_wind
                |--v_component_of_wind
                |--specific_humidity
                |--geopotential
            ...
  1. Ensemble Weather Forecast Post-Processing. This experiment employs the ENS-10 benchmark for global ensemble forecast post-processing, which pairs 10-member ensemble prediction (48-hour lead time) from the ECMWF Integrated Forecasting System (IFS) with ERA5 reanalysis targets at 0.5° resolution. Download the ENS10 and ERA5 data and follow the instructions provided by ENS-10. Excute data process script in ENS-10 (refer to ens10/baselines/utils and ens10/EFI) to process ERA5 and ens data and extract the EFI for additional variable U10 and V10. Please follow dataset structure below:
    |--WeatherPEFT
        |--datasets
            |--post-process
                |--process
                    |--efi_T2M.nc
                    |--efi_T850.nc
                    |--efi_U10M.nc
                    |--efi_V10M.nc
                    |--efi_Z500.nc
                    |--ENS10_sfc_mean_normalized.nc
                    |--ENS10_sfc_std_normalized.nc
                    |--ENS10_sfc_mean.nc
                    |--ENS10_sfc_std.nc
                    |--ENS10_pl_mean_500_normalized.nc
                    |--ENS10_pl_mean_500.nc
                    |--ENS10_pl_std_850_normalized.nc
                    |--ENS10_pl_std_850.nc
                    |--ERA5.nc
            ...
  1. Regional Precipitation Forecasting. In this experiment, we evaluate WeatherPEFT for regional precipitation forecasting, where the task is to predict the future six-hour accumulation of total precipitation (TP-6hr) based on the current weather conditions in the target region. The data is based on 12-hour interval 0.25° ERA5 data with 2010–2019 serving as the training set and 2020 as the test set, which can be downloaded from ECMWF. Process the hourly precipitation data into 6-hour cumulative precipitation data (refer to WeatherPEFT/dataset/process_tp_6hr.py) or directly download all this data from Weatherbench2. Process the data to package all the variable into one at each time point over the specific region (refer to WeatherPEFT/dataset/process.py). Please follow dataset structure below:
    |--WeatherPEFT
        |--datasets
            |--0.25
                |--china_organized
                    |--2010010100.npy
                    |--2010010112.npy
                    |--2010010200.npy
                    |--2010010212.npy
                    ...
                    |--2020123112.npy
            ...

Execution Instructions

Pretraining Checkponit

  • Download the pretraind-trained model weights of Aurora and and put them in WeatherPEFT/aurora-0.25-pretrained.ckpt

Fine-tuning

We provide three executable script in three folder (downscaling, postprocessing, and precipitation) for three corresponding downstream tasks.

For examble, using

cd downscaling
sh script_run_downscaling.sh

for downscaling task.

Acknowledgement

We thank the developers and authors of Aurora for releasing their pre-trained model and VideoMAE for releasing their helpful training framework.

Citation

If you use WeatherPEFT in your research or wish to refer to the baseline results, please use the following BibTeX entry.

@inproceedings{
cao2026taskadaptive,
title={Task-Adaptive Parameter-Efficient Fine-Tuning for Weather Foundation Models},
author={Shilei Cao and Hehai Lin and Jiashun Cheng and Yang Liu and Guowen Li and Xuehe Wang and Juepeng Zheng and Haoyuan Liang and Meng Jin and Chengwei Qin and Hong Cheng and Haohuan Fu},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=eFExhM3tKr}
}

About

[ICLR2026] Task-Adaptive Parameter-Efficient Fine-Tuning for Weather Foundation Models

Resources

Stars

8 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages