| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import torch |
| import torch.nn as nn |
|
|
| from model.dgmr_official.common import ContextConditioningStack, LatentConditioningStack |
| from model.dgmr_official.discriminators import Discriminator |
| from model.dgmr_official.generators import Generator, Sampler |
| from model.dgmr_official.losses import GridCellLoss |
|
|
|
|
| def weight_fn(y, precip_weight_cap=24.0): |
| """ |
| Weight function for the grid cell loss: w(y) = max(y + 1, cap). |
| """ |
| return torch.max(y + 1, torch.tensor(precip_weight_cap, device=y.device)) |
|
|
|
|
| class DGMR(nn.Module): |
| """ |
| Config-driven DGMR wrapper (generator + discriminator). |
| |
| The generator is a conditional GAN generator that takes ``num_context`` |
| observed radar frames of shape [B, T, C, H, W] and produces |
| ``forecast_steps`` future frames of the same spatial size. The |
| discriminator scores full sequences (context + forecast) spatially and |
| temporally; during GAN training the hinge losses plus the grid-cell |
| regularizer are applied (see ``dgmr_official/losses.py`` and the paper). |
| |
| Args: |
| forecast_steps: Number of frames to predict in the future (paper: 18). |
| num_context: Number of input/context frames (paper: 4). |
| input_channels: Number of channels per frame (paper: 1, radar). |
| output_shape: Spatial size of the frames; must be divisible by 32 |
| (paper: 256). Discriminators additionally need >= 128 px. |
| conv_type: Convolution flavour used by the conditioning stack, |
| one of "standard" / "coord" / "3d". |
| latent_channels / context_channels: DGMR architecture sizes |
| (paper: 768 / 384). |
| generation_steps: Number of Monte-Carlo generator samples used when |
| computing the grid-cell regularizer during training (paper: 6). |
| grid_lambda: Weight of the grid-cell regularizer (paper: 20). |
| precip_weight_cap: Ceiling for the grid-cell weight function (paper: 24). |
| """ |
|
|
| def __init__( |
| self, |
| forecast_steps: int = 18, |
| num_context: int = 4, |
| input_channels: int = 1, |
| output_shape: int = 256, |
| conv_type: str = "standard", |
| latent_channels: int = 768, |
| context_channels: int = 384, |
| generation_steps: int = 6, |
| grid_lambda: float = 20.0, |
| precip_weight_cap: float = 24.0, |
| ): |
| super().__init__() |
| self.forecast_steps = int(forecast_steps) |
| self.num_context = int(num_context) |
| self.input_channels = int(input_channels) |
| self.output_shape = int(output_shape) |
| self.conv_type = conv_type |
| self.latent_channels = int(latent_channels) |
| self.context_channels = int(context_channels) |
| self.generation_steps = int(generation_steps) |
| self.grid_lambda = float(grid_lambda) |
| self.precip_weight_cap = float(precip_weight_cap) |
|
|
| self.conditioning_stack = ContextConditioningStack( |
| input_channels=self.input_channels, |
| conv_type=self.conv_type, |
| output_channels=self.context_channels, |
| ) |
| self.latent_stack = LatentConditioningStack( |
| shape=( |
| 8 * self.input_channels, |
| self.output_shape // 32, |
| self.output_shape // 32, |
| ), |
| output_channels=self.latent_channels, |
| ) |
| self.sampler = Sampler( |
| forecast_steps=self.forecast_steps, |
| latent_channels=self.latent_channels, |
| context_channels=self.context_channels, |
| ) |
| self.generator = Generator(self.conditioning_stack, self.latent_stack, self.sampler) |
| self.discriminator = Discriminator(self.input_channels) |
| self.grid_regularizer = GridCellLoss( |
| weight_fn=weight_fn, precip_weight_cap=self.precip_weight_cap |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ |
| Args: |
| x: Observed radar frames, shape [batch, num_context, C, H, W]. |
| Returns: |
| Forecast frames, shape [batch, forecast_steps, C, H, W]. |
| """ |
| return self.generator(x) |
|
|
