metnet3 pytorch

MetNet-3 - Pytorch

Implementação do MetNet 3, modelo meteorológico neural SOTA do Google Deepmind, em Pytorch

A arquitetura do modelo é bastante normal. É basicamente uma U-net com um transformador de visão específico de bom desempenho. A coisa mais interessante sobre o artigo pode acabar sendo a escala de perdas na seção 4.3.2

Apreciação

• StabilityAI, programa de subsídios de IA de código aberto A16Z e ? Huggingface pelos generosos patrocínios, bem como aos meus outros patrocinadores, por me proporcionarem a independência para abrir o código-fonte da atual pesquisa de inteligência artificial

Instalar

$ pip install metnet3-pytorch

Uso

 import torch
from metnet3_pytorch import MetNet3

metnet3 = MetNet3 (
    dim = 512 ,
    num_lead_times = 722 ,
    lead_time_embed_dim = 32 ,
    input_spatial_size = 624 ,
    attn_dim_head = 8 ,
    hrrr_channels = 617 ,
    input_2496_channels = 2 + 14 + 1 + 2 + 20 ,
    input_4996_channels = 16 + 1 ,
    precipitation_target_bins = dict (
        mrms_rate = 512 ,
        mrms_accumulation = 512 ,
    ),
    surface_target_bins = dict (
        omo_temperature = 256 ,
        omo_dew_point = 256 ,
        omo_wind_speed = 256 ,
        omo_wind_component_x = 256 ,
        omo_wind_component_y = 256 ,
        omo_wind_direction = 180
    ),
    hrrr_loss_weight = 10 ,
    hrrr_norm_strategy = 'sync_batchnorm' ,  # this would use a sync batchnorm to normalize the input hrrr and target, without having to precalculate the mean and variance of the hrrr dataset per channel
    hrrr_norm_statistics = None             # you can also also set `hrrr_norm_strategy = "precalculated"` and pass in the mean and variance as shape `(2, 617)` through this keyword argument
)

# inputs

lead_times = torch . randint ( 0 , 722 , ( 2 ,))
hrrr_input_2496 = torch . randn (( 2 , 617 , 624 , 624 ))
hrrr_stale_state = torch . randn (( 2 , 1 , 624 , 624 ))
input_2496 = torch . randn (( 2 , 39 , 624 , 624 ))
input_4996 = torch . randn (( 2 , 17 , 624 , 624 ))

# targets

precipitation_targets = dict (
    mrms_rate = torch . randint ( 0 , 512 , ( 2 , 512 , 512 )),
    mrms_accumulation = torch . randint ( 0 , 512 , ( 2 , 512 , 512 )),
)

surface_targets = dict (
    omo_temperature = torch . randint ( 0 , 256 , ( 2 , 128 , 128 )),
    omo_dew_point = torch . randint ( 0 , 256 , ( 2 , 128 , 128 )),
    omo_wind_speed = torch . randint ( 0 , 256 , ( 2 , 128 , 128 )),
    omo_wind_component_x = torch . randint ( 0 , 256 , ( 2 , 128 , 128 )),
    omo_wind_component_y = torch . randint ( 0 , 256 , ( 2 , 128 , 128 )),
    omo_wind_direction = torch . randint ( 0 , 180 , ( 2 , 128 , 128 ))
)

hrrr_target = torch . randn ( 2 , 617 , 128 , 128 )

total_loss , loss_breakdown = metnet3 (
    lead_times = lead_times ,
    hrrr_input_2496 = hrrr_input_2496 ,
    hrrr_stale_state = hrrr_stale_state ,
    input_2496 = input_2496 ,
    input_4996 = input_4996 ,
    precipitation_targets = precipitation_targets ,
    surface_targets = surface_targets ,
    hrrr_target = hrrr_target
)

total_loss . backward ()

# after much training from above, you can predict as follows

metnet3 . eval ()

surface_preds , hrrr_pred , precipitation_preds = metnet3 (
    lead_times = lead_times ,
    hrrr_input_2496 = hrrr_input_2496 ,
    hrrr_stale_state = hrrr_stale_state ,
    input_2496 = input_2496 ,
    input_4996 = input_4996 ,
)


# Dict[str, Tensor], Tensor, Dict[str, Tensor] 

Pendência

• descobrir toda a entropia cruzada e perdas de MSE

• normalização de tratamento automático em todos os canais do HRRR, rastreando uma média contínua e variação de alvos durante o treinamento (usando sincronização batchnorm como hack)

• permitir que o pesquisador passe suas próprias variáveis ​​de normalização para HRRR

• construir todas as entradas de acordo com as especificações, também garantir que a entrada hrrr esteja normalizada, oferecer opção para desnormalizar as previsões hrrr

• certifique-se de que o modelo possa ser facilmente salvo e carregado, com diferentes maneiras de lidar com a norma hrrr

• descobrir a incorporação topológica, consulte um pesquisador de clima neural

Citações

 @article { Andrychowicz2023DeepLF ,
    title   = { Deep Learning for Day Forecasts from Sparse Observations } ,
    author  = { Marcin Andrychowicz and Lasse Espeholt and Di Li and Samier Merchant and Alexander Merose and Fred Zyda and Shreya Agrawal and Nal Kalchbrenner } ,
    journal = { ArXiv } ,
    year    = { 2023 } ,
    volume  = { abs/2306.06079 } ,
    url     = { https://api.semanticscholar.org/CorpusID:259129311 }
}

 @inproceedings { ElNouby2021XCiTCI ,
    title   = { XCiT: Cross-Covariance Image Transformers } ,
    author  = { Alaaeldin El-Nouby and Hugo Touvron and Mathilde Caron and Piotr Bojanowski and Matthijs Douze and Armand Joulin and Ivan Laptev and Natalia Neverova and Gabriel Synnaeve and Jakob Verbeek and Herv{'e} J{'e}gou } ,
    booktitle = { Neural Information Processing Systems } ,
    year    = { 2021 } ,
    url     = { https://api.semanticscholar.org/CorpusID:235458262 }
}