### Install MetNet-3

Source: https://github.com/lucidrains/metnet3-pytorch/blob/main/README.md

Install the package via pip.

```bash
$ pip install metnet3-pytorch
```

--------------------------------

### Complete MetNet-3 Training Loop

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

A full example of configuring the MetNet-3 model, setting up an optimizer, and executing a training loop with simulated data.

```python
import torch
import torch.optim as optim
from metnet3_pytorch import MetNet3

# Model configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
metnet3 = MetNet3(
    dim = 256,                          # Reduced for example
    num_lead_times = 722,
    lead_time_embed_dim = 32,
    input_spatial_size = 624,
    attn_depth = 4,                     # Reduced for example
    attn_dim_head = 8,
    attn_heads = 8,
    hrrr_channels = 64,                 # Reduced for example
    input_2496_channels = 39,
    input_4996_channels = 17,
    precipitation_target_bins = dict(
        mrms_rate = 128,
        mrms_accumulation = 128,
    ),
    surface_target_bins = dict(
        omo_temperature = 64,
        omo_dew_point = 64,
        omo_wind_speed = 64,
        omo_wind_component_x = 64,
        omo_wind_component_y = 64,
        omo_wind_direction = 90
    ),
    hrrr_loss_weight = 10,
    hrrr_norm_strategy = 'sync_batchnorm',
).to(device)

# Optimizer
optimizer = optim.AdamW(metnet3.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)

# Training loop
num_epochs = 10
batch_size = 2

for epoch in range(num_epochs):
    metnet3.train()

    # Simulated batch (replace with DataLoader)
    lead_times = torch.randint(0, 722, (batch_size,)).to(device)
    hrrr_input_2496 = torch.randn((batch_size, 64, 624, 624)).to(device)
    hrrr_stale_state = torch.randn((batch_size, 1, 624, 624)).to(device)
    input_2496 = torch.randn((batch_size, 39, 624, 624)).to(device)
    input_4996 = torch.randn((batch_size, 17, 624, 624)).to(device)

    precipitation_targets = dict(
        mrms_rate = torch.randint(0, 128, (batch_size, 512, 512)).to(device),
        mrms_accumulation = torch.randint(0, 128, (batch_size, 512, 512)).to(device),
    )
    surface_targets = dict(
        omo_temperature = torch.randint(0, 64, (batch_size, 128, 128)).to(device),
        omo_dew_point = torch.randint(0, 64, (batch_size, 128, 128)).to(device),
        omo_wind_speed = torch.randint(0, 64, (batch_size, 128, 128)).to(device),
        omo_wind_component_x = torch.randint(0, 64, (batch_size, 128, 128)).to(device),
        omo_wind_component_y = torch.randint(0, 64, (batch_size, 128, 128)).to(device),
        omo_wind_direction = torch.randint(0, 90, (batch_size, 128, 128)).to(device)
    )
    hrrr_target = torch.randn(batch_size, 64, 128, 128).to(device)

    # Forward pass
    optimizer.zero_grad()
    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
    )

    # Backward pass
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(metnet3.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()

    print(f"Epoch {epoch+1}/{num_epochs} - Loss: {total_loss.item():.4f}, "
          f"HRRR: {loss_breakdown.hrrr.item():.4f}, LR: {scheduler.get_last_lr()[0]:.6f}")

# Save trained model
metnet3.save('./metnet3_trained.pt')
print("Training complete!")
```

--------------------------------

### Initialize and Use ResnetBlocks

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Demonstrates how to instantiate ResnetBlocks with and without conditioning, and perform a forward pass.

```python
resnet_blocks = ResnetBlocks(
    dim = 512,              # Output dimension
    dim_in = 256,           # Input dimension (optional, defaults to dim)
    depth = 2,              # Number of residual blocks
    cond_dim = 32           # Conditioning dimension
)

# Forward pass
batch_size = 2
height, width = 64, 64
x = torch.randn(batch_size, 256, height, width)
cond = torch.randn(batch_size, 32)

output = resnet_blocks(x, cond=cond)
print(f"ResnetBlocks output shape: {output.shape}")  # (batch, 512, 64, 64)

# Without conditioning
resnet_blocks_no_cond = ResnetBlocks(
    dim = 512,
    dim_in = 256,
    depth = 3,
    cond_dim = None
)
output_no_cond = resnet_blocks_no_cond(x)
print(f"Output without conditioning: {output_no_cond.shape}")
```

--------------------------------

### Initialize and Train MetNet-3

Source: https://github.com/lucidrains/metnet3-pytorch/blob/main/README.md

Configure the MetNet-3 model with specific input channels and target bins, then perform a forward pass and backward propagation.

```python
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]
```

--------------------------------

### Initialize MetNet3 Model

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Initialize the MetNet3 model with a full configuration, specifying dimensions, channel counts, and target bins.

```python
import torch
from metnet3_pytorch import MetNet3

# Initialize MetNet3 model with full configuration
metnet3 = MetNet3(
    dim = 512,                          # Base model dimension
    num_lead_times = 722,               # Number of forecast lead times
    lead_time_embed_dim = 32,           # Lead time embedding dimension
    input_spatial_size = 624,           # Spatial size of input tensors
    attn_depth = 12,                    # Depth of MaxViT attention layers
    attn_dim_head = 8,                  # Dimension per attention head
    attn_heads = 32,                    # Number of attention heads
    attn_dropout = 0.1,                 # Attention dropout rate
    vit_window_size = 8,                # Window size for MaxViT
    vit_mbconv_expansion_rate = 4,      # MBConv expansion rate
    vit_mbconv_shrinkage_rate = 0.25,   # MBConv shrinkage rate
    hrrr_channels = 617,                # Number of HRRR channels
    input_2496_channels = 2 + 14 + 1 + 2 + 20,  # Channels at 2496m resolution (39 total)
    input_4996_channels = 16 + 1,       # Channels at 4996m resolution (17 total)
    surface_and_hrrr_target_spatial_size = 128,  # Output spatial size
    precipitation_target_bins = dict(
        mrms_rate = 512,                # Bins for precipitation rate
        mrms_accumulation = 512,        # Bins for precipitation accumulation
    ),
    surface_target_bins = dict(
        omo_temperature = 256,          # Bins for temperature
        omo_dew_point = 256,            # Bins for dew point
        omo_wind_speed = 256,           # Bins for wind speed
        omo_wind_component_x = 256,     # Bins for x wind component
        omo_wind_component_y = 256,     # Bins for y wind component
        omo_wind_direction = 180        # Bins for wind direction
    ),
    hrrr_loss_weight = 10,              # Weight for HRRR MSE loss
    hrrr_norm_strategy = 'sync_batchnorm',  # Normalization strategy
    hrrr_norm_statistics = None         # Optional precalculated statistics
)

print(f"Model initialized with {sum(p.numel() for p in metnet3.parameters()):,} parameters")
```

--------------------------------

### Initialize and Use MaxViT Component

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

The MaxViT module provides multi-axis attention. Ensure spatial dimensions are divisible by the configured window size.

```python
import torch
from metnet3_pytorch.metnet3_pytorch import MaxViT

# Initialize MaxViT independently
maxvit = MaxViT(
    dim = 512,                          # Base dimension
    depth = 12,                         # Number of transformer layers
    cond_dim = 32,                      # Conditioning dimension (lead time embedding)
    heads = 32,                         # Number of attention heads
    dim_head = 32,                      # Dimension per head
    window_size = 8,                    # Window size for local attention
    mbconv_expansion_rate = 4,          # MBConv expansion ratio
    mbconv_shrinkage_rate = 0.25,       # Squeeze-excitation shrinkage
    dropout = 0.1,                      # Dropout rate
    num_register_tokens = 4             # Number of register tokens
)

# Forward pass
batch_size = 2
spatial_size = 312  # Must be divisible by window_size
x = torch.randn(batch_size, 512, spatial_size, spatial_size)
cond = torch.randn(batch_size, 32)  # Conditioning vector (e.g., lead time embedding)

output = maxvit(x, cond=cond)
print(f"MaxViT output shape: {output.shape}")  # (batch, dim * 2^(num_stages-1), spatial/2^num_stages, ...)
```

--------------------------------

### Initialize ResnetBlocks

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

The ResnetBlocks module provides residual network layers with optional conditioning support.

```python
import torch
from metnet3_pytorch.metnet3_pytorch import ResnetBlocks
```

--------------------------------

### Training with Targets and Backpropagation

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Prepare targets for training and perform a forward pass. The model returns total loss and a breakdown of individual losses. Backpropagate the total loss and access individual loss components.

```python
precipitation_targets = dict(
    mrms_rate = torch.randint(0, 512, (batch_size, 512, 512)),
    mrms_accumulation = torch.randint(0, 512, (batch_size, 512, 512)),
)

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

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

# Forward pass with targets returns loss
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
)

# Backpropagation
total_loss.backward()

# Access individual losses
print(f"Total Loss: {total_loss.item():.4f}")
print(f"HRRR Loss: {loss_breakdown.hrrr.item():.4f}")
print(f"Surface Losses: {[(k, v.item()) for k, v in loss_breakdown.surface.items()]}")
print(f"Precipitation Losses: {[(k, v.item()) for k, v in loss_breakdown.precipitation.items()]}")
```

--------------------------------

### Save and Load MetNet-3 Checkpoints

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Persist model state and configuration to disk, with support for restoring via existing instances or direct class initialization.

```python
import torch
from metnet3_pytorch import MetNet3

# Initialize and train model
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 = 39,
    input_4996_channels = 17,
    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
    ),
)

# Save model (includes config and state_dict)
metnet3.save('./metnet3_checkpoint.pt', overwrite=True)
print("Model saved to ./metnet3_checkpoint.pt")

# Method 1: Load into existing model
metnet3_loaded = 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 = 39,
    input_4996_channels = 17,
)
metnet3_loaded.load('./metnet3_checkpoint.pt', strict=True)
print("Model loaded via load() method")

# Method 2: Initialize and load from checkpoint (recommended)
# Automatically restores all config parameters from saved checkpoint
metnet3_restored = MetNet3.init_and_load_from('./metnet3_checkpoint.pt', strict=True)
print("Model restored via init_and_load_from() class method")
```

--------------------------------

### HRRR Normalization: Precalculated Statistics

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Initialize MetNet3 with `hrrr_norm_strategy = 'precalculated'` and provide `hrrr_norm_statistics` containing precalculated mean and variance for HRRR data.

```python
import torch
from metnet3_pytorch import MetNet3

# Strategy 2: Precalculated mean/variance from dataset
hrrr_channels = 617
precalculated_stats = torch.randn(2, hrrr_channels)  # Shape: (2, channels) for mean and variance
precalculated_stats[1] = precalculated_stats[1].abs() + 0.1  # Ensure positive variance

metnet3_precalc = MetNet3(
    dim = 512,
    num_lead_times = 722,
    lead_time_embed_dim = 32,
    input_spatial_size = 624,
    attn_dim_head = 8,
    hrrr_channels = hrrr_channels,
    input_2496_channels = 39,
    input_4996_channels = 17,
    hrrr_norm_strategy = 'precalculated',
    hrrr_norm_statistics = precalculated_stats,  # Pass precalculated mean/variance
)
```

--------------------------------

### Configure Sync BatchNorm for Distributed Training

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Use the sync_batchnorm strategy when training across multiple devices to ensure consistent normalization.

```python
metnet3_syncbn = 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 = 39,
    input_4996_channels = 17,
    hrrr_norm_strategy = 'sync_batchnorm',  # Uses SyncBatchNorm for distributed training
)

print("Normalization strategies configured successfully")
```

--------------------------------

### HRRR Normalization: No Normalization

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Initialize MetNet3 with `hrrr_norm_strategy = 'none'` to disable any normalization for HRRR data.

```python
import torch
from metnet3_pytorch import MetNet3

# Strategy 1: No normalization
metnet3_no_norm = 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 = 39,
    input_4996_channels = 17,
    hrrr_norm_strategy = 'none',  # No normalization applied
)
```

--------------------------------

### Train MetNet3 Model

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

Prepare input tensors and targets for training the MetNet3 model. The model returns the total loss and a breakdown of individual losses.

```python
import torch
from metnet3_pytorch import MetNet3

# Initialize model
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 = 39,
    input_4996_channels = 17,
    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',
)

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

```

--------------------------------

### Implement XCAttention for Efficient Computation

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

XCAttention uses cross-covariance for linear attention. It supports optional conditioning vectors.

```python
import torch
from metnet3_pytorch.metnet3_pytorch import XCAttention

# Initialize XCAttention
xc_attn = XCAttention(
    dim = 512,              # Input/output dimension
    cond_dim = 32,          # Optional conditioning dimension
    dim_head = 32,          # Dimension per head
    heads = 8,              # Number of attention heads
    scale = 8,              # Attention scale factor
    dropout = 0.1           # Dropout rate
)

# Forward pass with conditioning
batch_size = 2
height, width = 64, 64
x = torch.randn(batch_size, 512, height, width)
cond = torch.randn(batch_size, 32)  # Lead time embedding

output = xc_attn(x, cond=cond)
print(f"XCAttention output shape: {output.shape}")  # (batch, 512, 64, 64)

# Forward pass without conditioning
xc_attn_no_cond = XCAttention(
    dim = 512,
    cond_dim = None,
    dim_head = 32,
    heads = 8,
)
output_no_cond = xc_attn_no_cond(x)
print(f"Output without conditioning: {output_no_cond.shape}")
```

--------------------------------

### Inference Mode without Targets

Source: https://context7.com/lucidrains/metnet3-pytorch/llms.txt

For inference, the model returns predictions without computing losses when targets are not provided. Ensure the model is in evaluation mode and use torch.no_grad() for efficiency.

```python
import torch
from metnet3_pytorch import MetNet3

# Initialize and load trained model
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 = 39,
    input_4996_channels = 17,
    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
    ),
)

# Switch to evaluation mode
metnet3.eval()

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

# Inference without targets returns predictions
with torch.no_grad():
    predictions = 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,
    )

# Unpack predictions (returns namedtuple: Predictions)
surface_preds = predictions.surface       # Dict[str, Tensor] - surface variable logits
hrrr_pred = predictions.hrrr              # Tensor - HRRR prediction
precipitation_preds = predictions.precipitation  # Dict[str, Tensor] - precipitation logits

# Access individual predictions
print(f"HRRR prediction shape: {hrrr_pred.shape}")  # (batch, 617, 128, 128)
print(f"Temperature logits shape: {surface_preds['omo_temperature'].shape}")  # (batch, 256, 128, 128)
print(f"Precipitation rate logits shape: {precipitation_preds['mrms_rate'].shape}")  # (batch, 512, 512, 512)

# Convert logits to class predictions
temperature_class = surface_preds['omo_temperature'].argmax(dim=1)
precipitation_rate_class = precipitation_preds['mrms_rate'].argmax(dim=1)
print(f"Temperature class predictions shape: {temperature_class.shape}")  # (batch, 128, 128)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.