DOCSEARCHbyAlgolia
Submit Docs

MetNet-3 PyTorch Implementation

repository·main·Indexed 42 days ago

https://github.com/lucidrains/metnet3-pytorch

A PyTorch implementation of Google DeepMind's MetNet-3 neural weather model. It utilizes a U-Net architecture combined with a vision transformer to predict weather variables, including precipitation and surface conditions, using HRRR, 2496, and 4996 input channels.

Tokens
1.2K
Snippets
4
Records
4
Agent score
4%

What's inside MetNet-3 PyTorch

  2. Train MetNet3 with loss calculation

    main

    During training, calling the metnet3 instance with both inputs and targets returns the total_loss and a loss_breakdown dictionary.

    Required Inputs:

    • lead_times: Tensor of lead times.
    • hrrr_input_2496: HRRR input tensor.
    • hrrr_stale_state: HRRR stale state tensor.
    • input_2496: Input 2496 tensor.
    • input_4996: Input 4996 tensor.

    Required Targets:

    • precipitation_targets: Dictionary of precipitation tensors.
    • surface_targets: Dictionary of surface variable tensors.
    • hrrr_target: HRRR target tensor.
    # Training step example
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()

  3. Predict weather with MetNet3

    main

    After training, set the model to evaluation mode (metnet3.eval()) and call it using only the input tensors. The model returns three outputs:

    1. surface_preds: A dictionary of predicted surface variables.
    2. hrrr_pred: The predicted HRRR tensor.
    3. precipitation_preds: A dictionary of predicted precipitation variables.
    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,
)

  4. Initialize the MetNet3 model

    main

    The MetNet3 class implements the MetNet 3 architecture (a U-Net with a vision transformer). When initializing, you must specify dimensions for the model, lead times, and input channels for various weather data sources (HRRR, 2496, and 4996).

    Key arguments include:

    • dim: Model dimension.
    • num_lead_times: Total number of lead times.
    • input_spatial_size: Spatial resolution of inputs.
    • hrrr_channels: Number of channels in the HRRR input.
    • precipitation_target_bins: A dictionary defining the number of bins for precipitation targets (e.g., mrms_rate, mrms_accumulation).
    • surface_target_bins: A dictionary defining the number of bins for surface targets (e.g., omo_temperature, omo_dew_point).
    • hrrr_norm_strategy: Strategy for HRRR normalization. Options include 'sync_batchnorm' or 'precalculated'.
    • hrrr_norm_statistics: If using 'precalculated', pass the mean and variance as a tensor of shape (2, hrrr_channels).
    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',
    hrrr_norm_statistics = None
)