File size: 7,988 Bytes
929e312 | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """Paper-version Samudra ConvNeXt U-Net implementation.
The layer layout follows the official ``samudra_om4_v1/model.yaml`` and the
official v1 ``blocks.py`` implementation. Inputs are channel-stacked tensors
with shape ``[batch, 158, lat, lon]`` for the full thermo-dynamic variant.
"""
from __future__ import annotations
from pathlib import Path
from typing import Iterable
import torch
from torch import nn
from torch.nn import functional as F
def circular_pad_width(x: torch.Tensor, padding: int) -> torch.Tensor:
"""Apply periodic longitude padding, including widths smaller than padding."""
if padding == 0:
return x
width = x.shape[-1]
repeats = (padding + width - 1) // width
tiled = x.repeat(1, 1, 1, repeats)
left = tiled[..., -padding:]
right = tiled[..., :padding]
return torch.cat((left, x, right), dim=-1)
def globe_pad(x: torch.Tensor, padding: int) -> torch.Tensor:
x = circular_pad_width(x, padding)
return F.pad(x, (0, 0, padding, padding), mode="constant")
class CappedGELU(nn.Module):
def __init__(self, cap_value: float = 10.0):
super().__init__()
self.gelu = nn.GELU()
# Keep this constant out of DDP buffers. DDP synchronizes buffers at
# every forward, which would mutate the autograd version counter during
# a recurrent multi-step rollout.
self.cap = float(cap_value)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.clamp(self.gelu(x), max=self.cap)
class ConvNeXtBlock(nn.Module):
"""Official v1 residual block with dilated 3x3 convolutions."""
def __init__(self, in_channels: int, out_channels: int, dilation: int):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.dilation = dilation
self.padding = dilation
self.skip = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, 1)
hidden = in_channels * 4
self.layers = nn.ModuleList(
[
nn.Conv2d(in_channels, hidden, 3, dilation=dilation),
nn.BatchNorm2d(hidden),
CappedGELU(),
nn.Conv2d(hidden, hidden, 3, dilation=dilation),
nn.BatchNorm2d(hidden),
CappedGELU(),
nn.Conv2d(hidden, out_channels, 1),
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
skip = self.skip(x)
for layer in self.layers:
if isinstance(layer, nn.Conv2d) and layer.kernel_size == (3, 3):
x = globe_pad(x, self.padding)
x = layer(x)
return skip + x
class PeriodicBilinearUpsample(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
width = x.shape[-1]
padded = F.pad(x, (1, 1, 0, 0), mode="circular")
result = F.interpolate(padded, scale_factor=2, mode="bilinear", align_corners=False)
return result[..., 2 : 2 + width * 2]
class SamudraUNet(nn.Module):
def __init__(self, in_channels: int, widths: Iterable[int], dilations: Iterable[int]):
super().__init__()
widths = list(widths)
dilations = list(dilations)
if len(widths) != 4 or len(dilations) != 4:
raise ValueError("paper v1 requires four widths and four dilation rates")
channels = [in_channels, *widths]
self.down_blocks = nn.ModuleList(
ConvNeXtBlock(channels[i], channels[i + 1], dilations[i]) for i in range(4)
)
self.pools = nn.ModuleList(nn.AvgPool2d(2) for _ in range(4))
self.middle = ConvNeXtBlock(widths[-1], widths[-1], dilations[-1])
self.first_up = PeriodicBilinearUpsample()
reversed_widths = list(reversed(widths))
reversed_dilations = list(reversed(dilations))
self.up_blocks = nn.ModuleList()
self.upsamples = nn.ModuleList()
current = widths[-1]
for index in range(3):
target = reversed_widths[index + 1]
self.up_blocks.append(ConvNeXtBlock(current, target, reversed_dilations[index]))
self.upsamples.append(PeriodicBilinearUpsample())
current = target
self.final = ConvNeXtBlock(current, widths[0], reversed_dilations[-1])
self.out_channels = widths[0]
def forward(self, x: torch.Tensor) -> torch.Tensor:
skips = []
for block, pool in zip(self.down_blocks, self.pools):
x = block(x)
skips.append(x)
x = pool(x)
x = self.middle(x)
x = self.first_up(x)
x = self._merge(x, skips[-1])
for block, up, skip in zip(self.up_blocks, self.upsamples, reversed(skips[:-1])):
x = block(x)
x = up(x)
x = self._merge(x, skip)
return self.final(x)
@staticmethod
def _merge(x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
height = min(x.shape[-2], skip.shape[-2])
width = min(x.shape[-1], skip.shape[-1])
x = x[..., :height, :width]
skip = skip[..., :height, :width]
if x.shape[1] != skip.shape[1]:
raise RuntimeError(f"skip channel mismatch: {x.shape[1]} != {skip.shape[1]}")
return x + skip
class Samudra(nn.Module):
"""Samudra v1 single-scale emulator.
``variant='thermo_dynamic'`` uses the paper's 158-to-154 interface.
``variant='thermo'`` uses 78 input channels and 154 output channels.
"""
def __init__(
self,
variant: str = "thermo_dynamic",
input_channels: int | None = None,
output_channels: int | None = None,
widths: Iterable[int] = (200, 250, 300, 400),
dilations: Iterable[int] = (1, 2, 4, 8),
):
super().__init__()
if variant not in {"thermo_dynamic", "thermo"}:
raise ValueError("variant must be thermo_dynamic or thermo")
default_input = 158 if variant == "thermo_dynamic" else 78
default_output = 154
self.variant = variant
self.input_channels = default_input if input_channels is None else input_channels
self.output_channels = default_output if output_channels is None else output_channels
self.unet = SamudraUNet(self.input_channels, widths, dilations)
self.decoder = nn.Conv2d(self.unet.out_channels, self.output_channels, 3)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.ndim != 4 or x.shape[1] != self.input_channels:
raise ValueError(f"expected [batch, {self.input_channels}, lat, lon], got {tuple(x.shape)}")
features = self.unet(x)
features = globe_pad(features, 1)
return self.decoder(features)
def load_official_checkpoint(self, path: str | Path, strict: bool = True) -> None:
checkpoint = torch.load(path, map_location="cpu")
if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
checkpoint = checkpoint["state_dict"]
if not isinstance(checkpoint, dict):
raise TypeError("official checkpoint must contain a state dictionary")
# Older local checkpoints stored the fixed GELU cap as a buffer.
checkpoint = {key: value for key, value in checkpoint.items() if not key.endswith(".cap")}
self.load_state_dict(checkpoint, strict=strict)
def build_model(config: dict) -> Samudra:
"""Build a model from the project's YAML model section."""
model_config = config.get("model", config)
model = Samudra(
variant=model_config.get("variant", "thermo_dynamic"),
input_channels=model_config.get("input_channels"),
output_channels=model_config.get("output_channels"),
widths=model_config.get("widths", (200, 250, 300, 400)),
dilations=model_config.get("dilations", (1, 2, 4, 8)),
)
checkpoint = model_config.get("checkpoint")
if checkpoint:
model.load_official_checkpoint(checkpoint)
return model
|