Aurora: Набор фундаментальных моделей ML от Microsoft Research for Science для ... - UPROGER

🌟 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"])