Aurora: Набор фундаментальных моделей ML от Microsoft Research for Science для прогнозирования атмосферных явлений

Архитектура базовой Aurora построена на 3D Swin Transformer. Модель обучалась более миллиона часов на корпусе данных о погоде и климате: прогнозах, параметрическом анализе и данных моделирования климата.

Длительное и обширное обучение позволило изучить модели общее представление динамики атмосферы, адаптироваться к широкому спектру задач прогнозирования: задачам с ограниченными входными данными, гетерогенными переменными и экстремальными событиями.

Aurora прошла двухэтапный процесс тонкой настройки. Предварительно настроенные веса сначала настраивались на короткое время прогнозирования (до 15 дней), а затем Aurora была интегрирована с LoRA комбинированных математических моделей длительного прогнозирования.

В наборе представлены три специализированных версии: одна для прогнозирования атмосферных явлений со средним разрешением, одна для прогнозирования с высоким разрешением и одна модель для прогнозирования загрязнения воздуха:

Aurora 0.1 - большая модель для прогнозов IFS HRES T0 c разрешением 0.1°;

Aurora 0.25 Pretrained - версия, обученная на обобщенном спектре параметров для использования в случае, если для целевого набора данных нету, например, для прогнозов на данных ERA5 с точностью 0.25°;

Aurora 0.25 Fine-Tuned - версия Aurora 0.25 Pretrained , дополнительно обученная на данных IFS HRES T0. Рекомендуется для прогнозирования на основе данных IFS HRES с разрешением 0.25°;

Aurora 0.25 Pretrained Small - уменьшенная версия Aurora 0.25 для процедур отладки. Не рекомендуется использовать для иных целей, кроме дебага;

Aurora air pollution forecasting - модель прогноза загрязнения воздуха, еще не опубликована, ее разработка активно ведется.

Модели семейства Aurora совместимы с данными прогнозных моделей ERA5 и HRES Европейского центра среднесрочных прогнозов погоды (ECMWF).

Установка и пример запуска c малой моделью и рандомными данными:

# Install with pip

pip install microsoft-aurora

#create a new venv & install the reqs:

virtualenv venv -p python3.10

source venv/bin/activate

make install

# Run the pretrained small model on random data:

from datetime import datetime

import torch

from aurora import AuroraSmall, Batch, Metadata

model = AuroraSmall()

model.load_checkpoint("microsoft/aurora", "aurora-0.25-small-pretrained.ckpt")

batch = Batch(

surf_vars={k: torch.randn(1, 2, 17, 32) for k in ("2t", "10u", "10v", "msl")},

static_vars={k: torch.randn(17, 32) for k in ("lsm", "z", "slt")},

atmos_vars={k: torch.randn(1, 2, 4, 17, 32) for k in ("z", "u", "v", "t", "q")},

metadata=Metadata(

lat=torch.linspace(90, -90, 17),

lon=torch.linspace(0, 360, 32 + 1)[:-1],

time=(datetime(2020, 6, 1, 12, 0),),

atmos_levels=(100, 250, 500, 850),

),

)

prediction = model.forward(batch)

print(prediction.surf_vars["2t"])

Лицензирование кода : MIT License.

Лицензирование моделей: CC-BY-NC-SA-4.0 License.

Набор моделей (https://huggingface.co/microsoft/aurora/tree/main)

Dev документация (https://microsoft.github.io/aurora/intro.html)

Arxiv (https://arxiv.org/pdf/2405.13063)

Github (https://github.com/microsoft/aurora) [ Stars: 120 | Issues: 1 | Forks: 12]

Источник: github.com

Комментарии: