File size: 5,712 Bytes
0ba2894 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | # coding=utf-8
#
# SPDX-License-Identifier: MIT
#
# Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
#
# This module is a thin configuration wrapper around the official Prithvi WxC
# model (Schmude et al., arXiv:2409.13598), taken verbatim from the
# NASA-IMPACT/Prithvi-WxC repository (commit 79dabfcd17abe77e2d5c696707c0164a04f2ec01,
# MIT License). The official implementation is vendored as
# ``prithvi_wxc_official.py``; only scaler construction and tensor-level call
# signatures are added here for YAML-driven usage.
import torch
import torch.nn as nn
from model.prithvi_wxc_official import PrithviWxC as _PrithviWxC
def _identity_scalers(n, device="cpu"):
return torch.zeros(n, device=device), torch.ones(n, device=device)
class PrithviWxC(nn.Module):
"""
Config-driven Prithvi WxC wrapper.
The official model is an encoder-decoder vision Transformer (Hiera + MaxViT,
alternating local/global attention) that maps an input state history to a
target state. This wrapper builds the official network with identity
scalers so that a small YAML config can be exercised; loading the 2.3B
checkpoints published on Hugging Face is possible by calling
``load_state_dict`` with the official checkpoint state (the buffer names,
e.g. ``input_scalers_mu``, are identical).
Args:
in_channels: number of input/output parameters.
input_size_time: number of input timestamps (paper uses 2).
in_channels_static: number of static input channels (>= 1).
n_lats_px / n_lons_px: grid size in pixels.
patch_size_px: patch/token size in pixels.
mask_unit_size_px: mask unit size in pixels.
mask_ratio_inputs: input masking ratio (0..1, 0 disables).
embed_dim / n_blocks_encoder / n_blocks_decoder / mlp_multiplier / n_heads:
transformer architecture hyper-parameters (paper: 2560 / 25 / 5 / 4 / 16).
dropout / drop_path / parameter_dropout: regularization ratios.
residual: one of "none" / "temporal" / "climate".
masking_mode: "global" / "local" / "both".
positional_encoding: "absolute" or "fourier".
encoder_shifting / decoder_shifting: Swin-style shifting.
"""
def __init__(
self,
in_channels=4,
input_size_time=2,
in_channels_static=4,
n_lats_px=32,
n_lons_px=64,
patch_size_px=(2, 2),
mask_unit_size_px=(8, 8),
mask_ratio_inputs=0.0,
embed_dim=16,
n_blocks_encoder=1,
n_blocks_decoder=1,
mlp_multiplier=4.0,
n_heads=4,
dropout=0.0,
drop_path=0.0,
parameter_dropout=0.0,
residual="none",
masking_mode="global",
positional_encoding="absolute",
encoder_shifting=False,
decoder_shifting=False,
):
super().__init__()
self.in_channels = int(in_channels)
self.input_size_time = int(input_size_time)
self.in_channels_static = int(in_channels_static)
self.n_lats_px = int(n_lats_px)
self.n_lons_px = int(n_lons_px)
self.residual = residual
self.positional_encoding = positional_encoding
mu_in, sigma_in = _identity_scalers(self.in_channels)
mu_s, sigma_s = _identity_scalers(self.in_channels_static)
self.model = _PrithviWxC(
in_channels=self.in_channels,
input_size_time=self.input_size_time,
in_channels_static=self.in_channels_static,
input_scalers_mu=mu_in,
input_scalers_sigma=sigma_in,
input_scalers_epsilon=1e-6,
static_input_scalers_mu=mu_s,
static_input_scalers_sigma=sigma_s,
static_input_scalers_epsilon=1e-6,
output_scalers=sigma_in,
n_lats_px=self.n_lats_px,
n_lons_px=self.n_lons_px,
patch_size_px=tuple(patch_size_px),
mask_unit_size_px=tuple(mask_unit_size_px),
mask_ratio_inputs=float(mask_ratio_inputs),
mask_ratio_targets=0.0,
embed_dim=int(embed_dim),
n_blocks_encoder=int(n_blocks_encoder),
n_blocks_decoder=int(n_blocks_decoder),
mlp_multiplier=int(mlp_multiplier),
n_heads=int(n_heads),
dropout=float(dropout),
drop_path=float(drop_path),
parameter_dropout=float(parameter_dropout),
residual=residual,
masking_mode=masking_mode,
positional_encoding=positional_encoding,
encoder_shifting=bool(encoder_shifting),
decoder_shifting=bool(decoder_shifting),
)
def forward(self, x, static, input_time=None, lead_time=None, climate=None):
"""
Args:
x: Tensor of shape [batch, time, parameter, lat, lon].
static: Tensor of shape [batch, static_channel, lat, lon].
input_time: Tensor of shape [batch] (optional, default zeros).
lead_time: Tensor of shape [batch] (optional, default zeros).
climate: Optional Tensor of shape [batch, parameter, lat, lon].
Returns:
Tensor of shape [batch, parameter, lat, lon].
"""
if input_time is None:
input_time = torch.zeros(x.shape[0], device=x.device)
if lead_time is None:
lead_time = torch.zeros(x.shape[0], device=x.device)
batch = {
"x": x,
"y": x[:, -1],
"static": static,
"input_time": input_time,
"lead_time": lead_time,
}
if climate is not None:
batch["climate"] = climate
return self.model(batch) |