diff --git a/olmoearth_pretrain/nn/naip_gan.py b/olmoearth_pretrain/nn/naip_gan.py new file mode 100644 index 000000000..f3e520595 --- /dev/null +++ b/olmoearth_pretrain/nn/naip_gan.py @@ -0,0 +1,794 @@ +"""Conditional pix2pix-GAN branch for predicting NAIP from the pooled embedding. + +The generator upsamples the encoder's pooled spatial embedding ``[B, H, W, D]`` +into a NAIP image, and a conditional PatchGAN discriminator distinguishes real +vs generated NAIP conditioned on the same pooled embedding. + +The generator lives inside the model (so it is trained by the main optimizer +alongside the encoder). The discriminator is held by the train module with its +own optimizer so its gradients never reach the encoder. +""" + +import logging +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor +from torch.distributed import DeviceMesh +from torch.distributed.fsdp import register_fsdp_forward_method + +from olmoearth_pretrain.config import Config +from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample, TokensAndMasks +from olmoearth_pretrain.nn.latent_mim import LatentMIM +from olmoearth_pretrain.nn.pooling import PoolingType, pool_unmasked_tokens + +logger = logging.getLogger(__name__) + + +def _num_groups(num_groups: int, channels: int) -> int: + """Return a valid GroupNorm group count that divides ``channels``.""" + return math.gcd(num_groups, channels) + + +def _apply_spectral_norm(module: nn.Module) -> None: + """Apply spectral normalization in-place to every Conv2d/Linear descendant. + + Uses :func:`torch.nn.utils.parametrizations.spectral_norm` (the parametrize + API), which registers a power-iteration parametrization on each layer's + ``weight``. Matches the ESRGAN / Real-ESRGAN ``UNetDiscriminatorSN`` practice + of spectrally normalizing all convolutions to bound the discriminator's + Lipschitz constant. + """ + from torch.nn.utils.parametrizations import spectral_norm + + for child_name, child in module.named_children(): + if isinstance(child, nn.Conv2d | nn.Linear): + setattr(module, child_name, spectral_norm(child)) + else: + _apply_spectral_norm(child) + + +class _ResBlock(nn.Module): + """A simple pre-activation residual block.""" + + def __init__(self, channels: int, num_groups: int = 8, norm: bool = True): + """Initialize the residual block. + + Args: + channels: Number of input/output channels. + num_groups: Target number of groups for GroupNorm. + norm: Whether to include GroupNorm before each activation. Set to + False for a spectrally-normalized discriminator: GroupNorm's + learned affine rescaling would otherwise undo the Lipschitz bound + spectral normalization is meant to impose (Real-ESRGAN's + ``UNetDiscriminatorSN`` is likewise norm-free). + """ + super().__init__() + groups = _num_groups(num_groups, channels) + norm1: nn.Module = nn.GroupNorm(groups, channels) if norm else nn.Identity() + norm2: nn.Module = nn.GroupNorm(groups, channels) if norm else nn.Identity() + self.block = nn.Sequential( + norm1, + nn.SiLU(), + nn.Conv2d(channels, channels, kernel_size=3, padding=1), + norm2, + nn.SiLU(), + nn.Conv2d(channels, channels, kernel_size=3, padding=1), + ) + + def forward(self, x: Tensor) -> Tensor: + """Apply the residual block.""" + return x + self.block(x) + + +class NaipGenerator(nn.Module): + """Generator that upsamples a pooled spatial embedding into a NAIP image.""" + + def __init__( + self, + embedding_size: int, + patch_size: int, + hidden_sizes: list[int], + out_channels: int = 4, + upsample_factor: int = 4, + num_res_blocks: int = 2, + num_groups: int = 8, + ): + """Initialize the generator. + + The total upsampling from the token grid is ``patch_size * upsample_factor``: + a learned per-patch unpatchify (``patch_size``) brings the token grid to the + base pixel resolution, then the conv trunk upsamples by ``upsample_factor``. + For NAIP set ``upsample_factor = naip_10.image_tile_size_factor``. + + To support a flexible / varying encoder patch size, pass the encoder + ``patch_size`` to :meth:`forward`: the pooled token grid is first bilinearly + resampled so each token spans the canonical ``self.patch_size`` base pixels + (keeping the physical extent fixed), and the single learned unpatchify then + always lands at the base resolution regardless of the encoder patch size. + + Args: + embedding_size: Channel dim ``D`` of the pooled embedding ``[B, H, W, D]``. + patch_size: Learned unpatchify block factor (the canonical patch + size); each (resampled) token is expanded into a + ``patch_size x patch_size`` block of features, e.g. 4 -> 40 m/px + tokens unpatchified to the 10 m/px base grid. + hidden_sizes: Per-stage channel widths, one for the base + (post-unpatchify) resolution followed by one for each upsampling + stage (length ``log2(upsample_factor) + 1``). Lets capacity be + concentrated at the coarse resolution, e.g. ``[256, 128, 128]``. + out_channels: Number of NAIP output bands (R, G, B, IR -> 4). + upsample_factor: Conv-trunk spatial upsampling factor (power of 2). + num_res_blocks: Residual blocks applied at the base resolution and + after each upsampling step. + num_groups: Target GroupNorm groups. + """ + super().__init__() + n_up = int(round(math.log2(upsample_factor))) + if 2**n_up != upsample_factor: + raise ValueError( + f"upsample_factor must be a power of 2, got {upsample_factor}" + ) + # Per-stage channel widths: one for the base (post-unpatchify) resolution + # plus one for each upsampling stage. + if len(hidden_sizes) != n_up + 1: + raise ValueError( + f"hidden_sizes must have length log2(upsample_factor) + 1 = " + f"{n_up + 1} (base resolution plus one per upsampling stage), " + f"got {len(hidden_sizes)}" + ) + channels = list(hidden_sizes) + self.patch_size = patch_size + self.hidden_size = channels[0] + self.upsample_factor = upsample_factor + self.out_channels = out_channels + + # Learned unpatchify: each token -> a patch_size x patch_size block of + # base-resolution features, i.e. the token grid is expanded to the base + # pixel grid. + self.unpatchify = nn.Linear( + embedding_size, patch_size * patch_size * channels[0] + ) + + layers: list[nn.Module] = [ + _ResBlock(channels[0], num_groups) for _ in range(num_res_blocks) + ] + for stage in range(n_up): + c_in = channels[stage] + c_out = channels[stage + 1] + # Sub-pixel (PixelShuffle) upsampling, ESRGAN-style: a conv expands to + # 4 * c_out channels and PixelShuffle(2) rearranges them into a 2x + # spatial block. This is a learnable upsampler (vs. fixed nearest) and + # tends to avoid the blockiness of nearest + conv. + layers.append(nn.Conv2d(c_in, c_out * 4, kernel_size=3, padding=1)) + layers.append(nn.PixelShuffle(2)) + layers.append(nn.GroupNorm(_num_groups(num_groups, c_out), c_out)) + layers.append(nn.SiLU()) + for _ in range(num_res_blocks): + layers.append(_ResBlock(c_out, num_groups)) + self.upsample = nn.Sequential(*layers) + + self.to_image = nn.Conv2d(channels[-1], out_channels, kernel_size=3, padding=1) + + def forward(self, pooled: Tensor, patch_size: int) -> Tensor: + """Generate a NAIP image. + + Args: + pooled: Pooled spatial embedding of shape ``[B, H, W, D]``. + patch_size: Encoder patch size for this batch. When it differs from + the unpatchify factor ``self.patch_size`` the token grid is + bilinearly resampled by ``patch_size / self.patch_size`` so each + token becomes canonical before the (fixed) unpatchify, keeping + the output resolution consistent across patch sizes. + + Returns: + NAIP image of shape ``[B, out_channels, Hc * U * upsample_factor, + Wc * U * upsample_factor]`` where ``U`` is the unpatchify factor + (``self.patch_size``) and ``(Hc, Wc)`` the (resampled) token grid. + """ + # Match the (possibly mixed-precision) param dtype; pooling can upcast to + # float32 even when the model runs in bf16 under FSDP mixed precision. + pooled = pooled.to(self.unpatchify.weight.dtype) + # Resample the token grid so each token spans the canonical + # ``self.patch_size`` base pixels (fixed extent), letting one learned + # unpatchify land at the base resolution for any encoder patch size. + h, w = pooled.shape[1], pooled.shape[2] + target = ( + max(1, round(h * patch_size / self.patch_size)), + max(1, round(w * patch_size / self.patch_size)), + ) + if target != (h, w): + pooled = ( + F.interpolate( + pooled.permute(0, 3, 1, 2), + size=target, + mode="bilinear", + align_corners=False, + ) + .permute(0, 2, 3, 1) + .contiguous() + ) + x = self.unpatchify(pooled) # [B, H, W, patch^2 * hidden] + x = rearrange( + x, + "b h w (ph pw c) -> b c (h ph) (w pw)", + ph=self.patch_size, + pw=self.patch_size, + c=self.hidden_size, + ) # [B, hidden, H * patch_size, W * patch_size] (base pixel grid) + x = self.upsample(x) + return self.to_image(x) + + +class NaipDiscriminator(nn.Module): + """Conditional PatchGAN discriminator over NAIP images. + + The image is downsampled with strided convs and adaptively pooled to the + fusion grid (the base 10 m/px grid in embedding mode, or the Sentinel-2 image + resolution in image mode), then combined with the projected condition before + a conv head produces per-patch real/fake logits. + + Because the condition here is a rich Sentinel-2-derived embedding (not a + downscaled NAIP image), the discriminator can be pushed to actually judge + image-vs-condition *consistency* rather than unconditional realism via: + + * post-unpatchify convs on the condition (``cond_embedding_channels``), + * a deeper reasoning head (``num_head_res_blocks``), and + * a projection-discriminator inner-product term (``use_projection``; + Miyato & Koyama, "cGANs with Projection Discriminator"), which only + rewards logits when image features align with the projected condition. + """ + + def __init__( + self, + embedding_size: int, + in_channels: int = 4, + image_strided_conv_channels: list[int] | None = None, + feature_channels: int = 256, + cond_embedding_channels: list[int] | None = None, + num_head_res_blocks: int = 0, + use_projection: bool = False, + num_groups: int = 8, + cond_mode: str = "embedding", + cond_in_channels: int = 4, + cond_image_pre_pool_channels: list[int] | None = None, + cond_image_post_pool_channels: list[int] | None = None, + num_convs_per_resolution: int = 0, + cond_unpatchify_factor: int = 1, + use_spectral_norm: bool = False, + ): + """Initialize the discriminator. + + Args: + embedding_size: Channel dim ``D`` of the conditioning embedding + (used only when ``cond_mode == 'embedding'``). + in_channels: Number of NAIP input bands. + image_strided_conv_channels: Per-conv output widths of the NAIP image + stack. The first entry is a stride-1 stem and each subsequent + entry a stride-2 conv, e.g. ``[64, 128, 256]`` is a stem to 64 + then two strided convs to 128 and 256. Defaults to + ``[64, 128, 256]``. + feature_channels: Fused feature width. A final conv projects the image + stack to this width, the condition is produced at this width, and + the head fuses the two (so its input is ``2 * feature_channels``). + cond_embedding_channels: Optional per-conv output widths of the + non-strided convs applied to the condition *after* the learned + unpatchify (a final 1x1 conv projects back to + ``feature_channels``), e.g. ``[256]`` gives + ``feature_channels -> 256 -> feature_channels``. ``None``/empty + applies no post-unpatchify convs. Only used when + ``cond_mode == 'embedding'``. + num_head_res_blocks: Number of residual blocks inserted into the + fusion head so it has depth to compare image features against the + condition (0 keeps the original shallow head). + use_projection: If True, add a projection-discriminator term: the + spatial inner product between the image features and the + projected condition is added to the logits, so the discriminator + cannot score well by ignoring the condition. + num_groups: Target GroupNorm groups for the head residual blocks. + cond_mode: How the condition is provided. ``'embedding'`` (default) + takes a pooled embedding ``[B, H, W, D]`` and projects it; + ``'image'`` takes a raw image temporal stack + ``[B, T, cond_in_channels, Hc, Wc]`` (the Sentinel-2 time series) + and embeds it (see ``cond_image_pre_pool_channels``). + cond_in_channels: Number of channels (bands) of the raw image + condition when ``cond_mode == 'image'``. + cond_image_pre_pool_channels: Per-conv output widths of the + non-strided convs applied to each timestep of the raw image + condition (at the image resolution) before the mean over time. + Defaults to ``[feature_channels]``. + cond_image_post_pool_channels: Per-conv output widths of the extra + non-strided convs applied after the mean over time (before the + final 1x1 projection to ``feature_channels``). Defaults to none. + num_convs_per_resolution: Number of stride-1 refinement convs (3x3 + + LeakyReLU) inserted after each conv of the NAIP image stack, at + that resolution (0 disables them). + cond_unpatchify_factor: In ``'embedding'`` mode, the factor ``f`` of + the learned unpatchify: tokens are first resampled so each spans + ``f`` canonical base pixels, then a single linear expands each + into an ``f x f`` block of features, mirroring the generator. + Fusion always happens at the base ``token_grid * patch_size`` + (10 m/px) grid regardless of ``f``; ``f`` only sets the + unpatchify block factor (set it to the generator ``patch_size``, + e.g. 4, to resample tokens to 40 m/px before a learned 4x4 + unpatchify to 10 m/px). + use_spectral_norm: If True, wrap every ``Conv2d``/``Linear`` in the + discriminator with spectral normalization (Miyato et al., + matching the ESRGAN / Real-ESRGAN ``UNetDiscriminatorSN``). This + bounds the discriminator's Lipschitz constant and is the main + stabilizer when the reconstruction (L1) anchor is weak or absent. + """ + super().__init__() + if cond_mode not in ("embedding", "image"): + raise ValueError(f"Unknown cond_mode: {cond_mode}") + if cond_unpatchify_factor < 1: + raise ValueError( + f"cond_unpatchify_factor must be >= 1, got {cond_unpatchify_factor}" + ) + self.cond_mode = cond_mode + self.use_projection = use_projection + self.cond_unpatchify_factor = cond_unpatchify_factor + self.feature_channels = feature_channels + + # Downsample the NAIP image: a stride-1 stem (channels[0]) then stride-2 + # convs (channels[1:]), each optionally followed by num_convs_per_resolution + # stride-1 refinement convs, and a final conv to feature_channels. + image_channels = image_strided_conv_channels or [64, 128, 256] + if not image_channels: + raise ValueError("image_strided_conv_channels must be non-empty") + image_layers: list[nn.Module] = [] + prev = in_channels + for i, width in enumerate(image_channels): + if i == 0: + image_layers.append( + nn.Conv2d(prev, width, kernel_size=3, stride=1, padding=1) + ) + else: + image_layers.append( + nn.Conv2d(prev, width, kernel_size=4, stride=2, padding=1) + ) + image_layers.append(nn.LeakyReLU(0.2, inplace=True)) + for _ in range(num_convs_per_resolution): + image_layers.append( + nn.Conv2d(width, width, kernel_size=3, stride=1, padding=1) + ) + image_layers.append(nn.LeakyReLU(0.2, inplace=True)) + prev = width + image_layers.append(nn.Conv2d(prev, feature_channels, kernel_size=3, padding=1)) + image_layers.append(nn.LeakyReLU(0.2, inplace=True)) + self.from_image = nn.Sequential(*image_layers) + + if cond_mode == "embedding": + # Learned unpatchify of the pooled embedding [B, H, W, D]: a single + # linear expands each token into an f x f block of features + # (f = cond_unpatchify_factor), mirroring the generator. + cond_out = ( + cond_unpatchify_factor * cond_unpatchify_factor * feature_channels + ) + self.cond_proj = nn.Linear(embedding_size, cond_out) + # Optional non-strided convs applied after the unpatchify, with a + # final 1x1 projection back to feature_channels. Identity when unset. + unpatchify_layers: list[nn.Module] = [] + prev = feature_channels + for width in cond_embedding_channels or []: + unpatchify_layers.append( + nn.Conv2d(prev, width, kernel_size=3, padding=1) + ) + unpatchify_layers.append(nn.LeakyReLU(0.2, inplace=True)) + prev = width + if cond_embedding_channels: + unpatchify_layers.append( + nn.Conv2d(prev, feature_channels, kernel_size=1) + ) + self.cond_post_unpatchify = nn.Sequential(*unpatchify_layers) + else: + # Raw image condition as a temporal stack [B, T, cond_in_channels, H, W]. + # Non-strided per-timestep convs keep it at the image resolution; after + # the mean over time, more non-strided convs and a 1x1 projection to + # feature_channels. The result is resampled to the fusion grid in + # forward (a no-op when it already matches). + pre_channels = cond_image_pre_pool_channels or [feature_channels] + pre_layers: list[nn.Module] = [] + prev = cond_in_channels + for width in pre_channels: + pre_layers.append(nn.Conv2d(prev, width, kernel_size=3, padding=1)) + pre_layers.append(nn.LeakyReLU(0.2, inplace=True)) + prev = width + self.cond_pre_pool = nn.Sequential(*pre_layers) + post_layers: list[nn.Module] = [] + for width in cond_image_post_pool_channels or []: + post_layers.append(nn.Conv2d(prev, width, kernel_size=3, padding=1)) + post_layers.append(nn.LeakyReLU(0.2, inplace=True)) + prev = width + post_layers.append(nn.Conv2d(prev, feature_channels, kernel_size=1)) + self.cond_post_pool = nn.Sequential(*post_layers) + + head_layers: list[nn.Module] = [ + nn.Conv2d( + self.feature_channels * 2, + self.feature_channels, + kernel_size=3, + padding=1, + ), + nn.LeakyReLU(0.2, inplace=True), + ] + for _ in range(num_head_res_blocks): + head_layers.append( + _ResBlock(self.feature_channels, num_groups, norm=not use_spectral_norm) + ) + head_layers.append( + nn.Conv2d(self.feature_channels, 1, kernel_size=3, padding=1) + ) + self.head = nn.Sequential(*head_layers) + + if use_spectral_norm: + _apply_spectral_norm(self) + + def forward( + self, + image: Tensor, + cond: Tensor, + patch_size: int, + cond_time_mask: Tensor | None = None, + ) -> Tensor: + """Classify each patch as real or fake. + + Args: + image: NAIP image of shape ``[B, in_channels, Hf, Wf]``. + cond: Conditioning input. When ``cond_mode == 'embedding'`` this is a + pooled embedding ``[B, H, W, D]`` (whose ``[H, W]`` is the token + grid); when ``cond_mode == 'image'`` this is a raw image temporal + stack ``[B, T, cond_in_channels, Hc, Wc]`` (the Sentinel-2 time + series at its native 10 m/px resolution). + patch_size: Encoder patch size for this batch. Only used in + ``'embedding'`` mode: the token grid is resampled by + ``patch_size / cond_unpatchify_factor`` so each token spans + ``cond_unpatchify_factor`` canonical base pixels before the + learned unpatchify, landing fusion at the base + ``token_grid * patch_size`` (10 m/px) grid regardless of patch + size. Ignored in ``'image'`` mode. + cond_time_mask: Optional ``[B, T]`` boolean mask of valid timesteps for + the ``'image'`` condition; the mean over time uses only valid + timesteps. When ``None`` all timesteps are averaged. + + Returns: + Per-patch logits. In ``'embedding'`` mode of shape + ``[B, 1, H * patch_size, W * patch_size]`` (the base 10 m/px grid); + in ``'image'`` mode of shape ``[B, 1, Hc, Wc]`` (the Sentinel-2 image + resolution). + """ + # The discriminator is kept in its own (fp32) precision, separate from the + # FSDP mixed-precision model, so cast inputs to its param dtype. Reference + # the first image conv, which exists regardless of the cond variant. + param_dtype = self.from_image[0].weight.dtype + image = image.to(param_dtype) + cond = cond.to(param_dtype) + feats = self.from_image(image) + f = self.cond_unpatchify_factor + if self.cond_mode == "embedding": + token_grid = (int(cond.shape[1]), int(cond.shape[2])) + # Resample tokens so each spans the canonical ``f`` base pixels, so the + # learned unpatchify lands at the base (10 m/px) grid for any encoder + # patch size. + canonical = ( + max(1, round(token_grid[0] * patch_size / f)), + max(1, round(token_grid[1] * patch_size / f)), + ) + if canonical != token_grid: + cond = ( + F.interpolate( + cond.permute(0, 3, 1, 2), + size=canonical, + mode="bilinear", + align_corners=False, + ) + .permute(0, 2, 3, 1) + .contiguous() + ) + token_grid = canonical + fusion_grid = (token_grid[0] * f, token_grid[1] * f) + feats = F.adaptive_avg_pool2d(feats, output_size=fusion_grid) + # Learned unpatchify: [B, H, W, f*f*C] -> [B, C, H*f, W*f]. + cond_feats = rearrange( + self.cond_proj(cond), + "b h w (ph pw c) -> b c (h ph) (w pw)", + ph=f, + pw=f, + c=self.feature_channels, + ).contiguous() + cond_feats = self.cond_post_unpatchify(cond_feats) + else: + # cond is a temporal stack [B, T, cond_in_channels, H, W] at the + # Sentinel-2 (10 m/px) resolution. Apply the per-timestep convs, mean + # over the (valid) timesteps, then the post-pool convs; fusion happens + # at that native image resolution. + b, t = cond.shape[0], cond.shape[1] + x = self.cond_pre_pool(cond.flatten(0, 1)).unflatten(0, (b, t)) + if cond_time_mask is not None: + m = cond_time_mask.to(x.dtype).reshape(b, t, 1, 1, 1) + x = (x * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0) + else: + x = x.mean(dim=1) + cond_feats = self.cond_post_pool(x) # [B, feature_channels, Hc, Wc] + feats = F.adaptive_avg_pool2d(feats, output_size=cond_feats.shape[-2:]) + combined = torch.cat([feats, cond_feats], dim=1) + logits = self.head(combined) + if self.use_projection: + # Projection-discriminator term: reward image/condition agreement. + proj = (feats * cond_feats).sum(dim=1, keepdim=True) + logits = logits + proj + return logits + + +def discriminator_adversarial_loss( + real_logits: Tensor, fake_logits: Tensor, loss_type: str = "hinge" +) -> Tensor: + """Adversarial loss for the discriminator (real should score high).""" + if loss_type == "hinge": + return F.relu(1.0 - real_logits).mean() + F.relu(1.0 + fake_logits).mean() + if loss_type == "bce": + return F.binary_cross_entropy_with_logits( + real_logits, torch.ones_like(real_logits) + ) + F.binary_cross_entropy_with_logits( + fake_logits, torch.zeros_like(fake_logits) + ) + raise ValueError(f"Unknown GAN loss type: {loss_type}") + + +def generator_adversarial_loss(fake_logits: Tensor, loss_type: str = "hinge") -> Tensor: + """Adversarial loss for the generator (fake should fool the discriminator).""" + if loss_type == "hinge": + return -fake_logits.mean() + if loss_type == "bce": + return F.binary_cross_entropy_with_logits( + fake_logits, torch.ones_like(fake_logits) + ) + raise ValueError(f"Unknown GAN loss type: {loss_type}") + + +@dataclass +class NaipGeneratorConfig(Config): + """Configuration for :class:`NaipGenerator`. + + Total upsampling from the token grid is ``patch_size * upsample_factor``; set + ``upsample_factor`` to the NAIP ``image_tile_size_factor`` so the output lands + at native NAIP resolution. + + Set ``hidden_sizes`` to a per-stage list of channel widths (base resolution + plus one per upsampling stage, length ``log2(upsample_factor) + 1``) to + concentrate capacity at the coarse resolution, e.g. ``[256, 128, 128]``. + + Pass the encoder ``patch_size`` to the generator's ``forward`` to support a + flexible encoder patch size: the token grid is resampled to the canonical + ``patch_size`` (unpatchify factor) before the unpatchify, so the output stays + consistent across patch sizes. + """ + + embedding_size: int + patch_size: int + hidden_sizes: list[int] + out_channels: int = 4 + upsample_factor: int = 4 + num_res_blocks: int = 2 + num_groups: int = 8 + + def build(self) -> NaipGenerator: + """Build the generator.""" + return NaipGenerator( + embedding_size=self.embedding_size, + patch_size=self.patch_size, + hidden_sizes=self.hidden_sizes, + out_channels=self.out_channels, + upsample_factor=self.upsample_factor, + num_res_blocks=self.num_res_blocks, + num_groups=self.num_groups, + ) + + +@dataclass +class NaipDiscriminatorConfig(Config): + """Configuration for :class:`NaipDiscriminator`. + + Set ``cond_embedding_channels``, ``num_head_res_blocks`` and ``use_projection`` + to make the discriminator reason more about the conditioning embedding (rather + than just the realism of the input image). + + Set ``cond_mode='image'`` (with ``cond_in_channels``) to condition on a raw + image (the full Sentinel-2 temporal stack) instead of a pooled embedding; + ``embedding_size`` is then unused. ``cond_image_pre_pool_channels`` are + non-strided convs applied to each timestep at the S2 resolution; + ``cond_image_post_pool_channels`` are non-strided convs applied after the mean + over time. In embedding mode ``cond_embedding_channels`` are non-strided convs + applied after the learned unpatchify of the condition. + + ``image_strided_conv_channels`` is the per-conv width list for the NAIP image + stack (first entry a stride-1 stem, the rest stride-2 convs); ``feature_channels`` + is the fused width the image and condition are projected to before the head. + + Set ``cond_unpatchify_factor`` to the generator ``patch_size`` so the condition + tokens are resampled to that (coarse) resolution and a learned unpatchify + expands them back to the base 10 m/px fusion grid, mirroring the generator + (embedding mode only; image mode fuses at the Sentinel-2 image resolution). + """ + + embedding_size: int + in_channels: int = 4 + image_strided_conv_channels: list[int] | None = None + feature_channels: int = 256 + cond_embedding_channels: list[int] | None = None + num_head_res_blocks: int = 0 + use_projection: bool = False + num_groups: int = 8 + cond_mode: str = "embedding" + cond_in_channels: int = 4 + cond_image_pre_pool_channels: list[int] | None = None + cond_image_post_pool_channels: list[int] | None = None + num_convs_per_resolution: int = 0 + cond_unpatchify_factor: int = 1 + use_spectral_norm: bool = False + + def build(self) -> NaipDiscriminator: + """Build the discriminator.""" + return NaipDiscriminator( + embedding_size=self.embedding_size, + in_channels=self.in_channels, + image_strided_conv_channels=self.image_strided_conv_channels, + feature_channels=self.feature_channels, + cond_embedding_channels=self.cond_embedding_channels, + num_head_res_blocks=self.num_head_res_blocks, + use_projection=self.use_projection, + num_groups=self.num_groups, + cond_mode=self.cond_mode, + cond_in_channels=self.cond_in_channels, + cond_image_pre_pool_channels=self.cond_image_pre_pool_channels, + cond_image_post_pool_channels=self.cond_image_post_pool_channels, + num_convs_per_resolution=self.num_convs_per_resolution, + cond_unpatchify_factor=self.cond_unpatchify_factor, + use_spectral_norm=self.use_spectral_norm, + ) + + +class NaipGanModel(LatentMIM): + """LatentMIM plus a NAIP generator on top of the pooled spatial embedding.""" + + def __init__( + self, + encoder: nn.Module, + decoder: nn.Module, + generator: nn.Module, + reconstructor: torch.nn.Module | None = None, + ): + """Initialize the model. + + Args: + encoder: The online encoder. + decoder: The latent-MIM decoder (predictor). + generator: The NAIP generator applied to the pooled embedding. + reconstructor: Optional MAE reconstructor. + """ + super().__init__(encoder=encoder, decoder=decoder, reconstructor=reconstructor) + self.generator = generator + + def forward( # type: ignore[override] + self, x: MaskedOlmoEarthSample, patch_size: int + ) -> tuple[ + TokensAndMasks, + TokensAndMasks, + torch.Tensor, + TokensAndMasks | None, + dict, + torch.Tensor, + torch.Tensor, + ]: + """Forward pass returning the standard LatentMIM outputs plus GAN tensors. + + Returns: + The five LatentMIM outputs, followed by the pooled spatial embedding + ``[B, H, W, D]`` and the generated NAIP image ``[B, C, Hf, Wf]``. + """ + latent, decoded, latent_projected_and_pooled, reconstructed, extra_metrics = ( + super().forward(x, patch_size) + ) + pooled = pool_unmasked_tokens( + latent, PoolingType.MEAN, spatial_pooling=True + ) # [B, H, W, D] + fake_naip = self.generator(pooled, patch_size) + return ( + latent, + decoded, + latent_projected_and_pooled, + reconstructed, + extra_metrics, + pooled, + fake_naip, + ) + + def apply_fsdp( + self, + dp_mesh: DeviceMesh | None = None, + param_dtype: torch.dtype | None = None, + reduce_dtype: torch.dtype = torch.float32, + prefetch_factor: int = 0, + ) -> None: + """Apply FSDP, additionally registering the online encoder's ``forward``. + + The base class registers the target encoder's ``forward`` as an FSDP + forward method so it can be called directly (it unshards/reshards around + the call). The discriminator's ``online_unmasked_pooled`` condition calls + ``self.encoder.forward`` directly too, so register it the same way; + otherwise its (sharded) params stay ``DTensor`` and mix with the plain + input tensor. + """ + super().apply_fsdp( + dp_mesh=dp_mesh, + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + prefetch_factor=prefetch_factor, + ) + register_fsdp_forward_method(self.encoder, "forward") + + +@dataclass +class NaipGanModelConfig(Config): + """Configuration for :class:`NaipGanModel`. + + Mirrors ``LatentMIMConfig`` (encoder + decoder + optional reconstructor) and + adds the generator. Kept as a standalone config (rather than subclassing + ``LatentMIMConfig``) so the ``generator_config`` field is required. + """ + + encoder_config: Config + decoder_config: Config + generator_config: Config + reconstructor_config: Config | None = None + + def validate(self) -> None: + """Validate encoder/decoder compatibility (mirrors LatentMIMConfig).""" + if ( + self.encoder_config.supported_modalities + != self.decoder_config.supported_modalities + ): + raise ValueError("Encoder and decoder must support the same modalities") + if ( + self.encoder_config.max_sequence_length + != self.decoder_config.max_sequence_length + ): + raise ValueError( + "Encoder and decoder must have the same max sequence length" + ) + encoder_output_size = ( + self.encoder_config.output_embedding_size + or self.encoder_config.embedding_size + ) + if encoder_output_size != self.decoder_config.encoder_embedding_size: + raise ValueError("Encoder embedding size must be consistent!") + if encoder_output_size != self.generator_config.embedding_size: + raise ValueError( + "Generator embedding_size must match the encoder output size " + f"({encoder_output_size} != {self.generator_config.embedding_size})" + ) + # The generator (and the discriminator, in the train module) resamples the + # token grid to the canonical unpatchify factor using the per-batch encoder + # patch size, so a fixed encoder patch size is no longer required. + + def build(self) -> NaipGanModel: + """Build the NAIP GAN model.""" + self.validate() + encoder = self.encoder_config.build() + decoder = self.decoder_config.build() + reconstructor = ( + self.reconstructor_config.build() + if self.reconstructor_config is not None + else None + ) + generator = self.generator_config.build() + return NaipGanModel( + encoder=encoder, + decoder=decoder, + generator=generator, + reconstructor=reconstructor, + ) diff --git a/olmoearth_pretrain/nn/perceptual.py b/olmoearth_pretrain/nn/perceptual.py new file mode 100644 index 000000000..c9fefdb91 --- /dev/null +++ b/olmoearth_pretrain/nn/perceptual.py @@ -0,0 +1,156 @@ +"""VGG perceptual (content) loss for the NAIP GAN branch. + +Mirrors the perceptual loss used by ESRGAN / Real-ESRGAN (and the satlas +super-resolution recipe): a frozen ImageNet VGG-19 feature extractor compares +generated and real images in feature space at several depths, using the +pre-activation (before-ReLU) conv outputs with the standard layer weights. This +supplies the mid-frequency content/structure gradient that an L1 pixel loss and +a weak adversarial term cannot provide on their own. +""" + +import logging +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +logger = logging.getLogger(__name__) + +# Indices into ``torchvision`` ``vgg19().features`` of the conv outputs +# *before* their ReLU, matching the ESRGAN layer names. +_VGG19_PRE_RELU_INDICES: dict[str, int] = { + "conv1_2": 2, + "conv2_2": 7, + "conv3_4": 16, + "conv4_4": 25, + "conv5_4": 34, +} + +# ESRGAN / satlas default per-layer weights. +_DEFAULT_LAYER_WEIGHTS: dict[str, float] = { + "conv1_2": 0.1, + "conv2_2": 0.1, + "conv3_4": 1.0, + "conv4_4": 1.0, + "conv5_4": 1.0, +} + +# ImageNet normalization statistics (VGG was trained on these). +_IMAGENET_MEAN = (0.485, 0.456, 0.406) +_IMAGENET_STD = (0.229, 0.224, 0.225) + + +class VGGPerceptualLoss(nn.Module): + """Frozen VGG-19 perceptual loss over pre-ReLU feature maps. + + The module is a pure loss: its parameters are frozen and excluded from every + optimizer, but gradients still flow through it into whatever produced the + ``fake`` image. Inputs are RGB images in ``[0, 1]`` and are ImageNet-normalized + internally. + """ + + def __init__( + self, + layer_weights: dict[str, float] | None = None, + criterion: str = "l1", + ): + """Initialize the perceptual loss. + + Args: + layer_weights: Mapping from VGG-19 layer name (a key of + ``_VGG19_PRE_RELU_INDICES``) to its weight in the summed loss. + Defaults to the ESRGAN weights. + criterion: Feature comparison criterion, ``"l1"`` or ``"l2"``. + """ + super().__init__() + from torchvision.models import VGG19_Weights, vgg19 + + weights = layer_weights or dict(_DEFAULT_LAYER_WEIGHTS) + unknown = set(weights) - set(_VGG19_PRE_RELU_INDICES) + if unknown: + raise ValueError( + f"Unknown VGG layer(s) {sorted(unknown)}; valid layers are " + f"{sorted(_VGG19_PRE_RELU_INDICES)}" + ) + if criterion not in ("l1", "l2"): + raise ValueError(f"criterion must be 'l1' or 'l2', got {criterion!r}") + self.layer_weights = weights + self.criterion = criterion + + # Keep only the feature layers up to the deepest one we need. + max_index = max(_VGG19_PRE_RELU_INDICES[name] for name in weights) + features = vgg19(weights=VGG19_Weights.IMAGENET1K_V1).features[: max_index + 1] + self.features = features.eval() + for param in self.features.parameters(): + param.requires_grad_(False) + + # ``[1, 3, 1, 1]`` normalization buffers so they follow ``.to(device)``. + self.register_buffer( + "mean", torch.tensor(_IMAGENET_MEAN).reshape(1, 3, 1, 1), persistent=False + ) + self.register_buffer( + "std", torch.tensor(_IMAGENET_STD).reshape(1, 3, 1, 1), persistent=False + ) + + # Map each captured feature-map index to its layer weight. + self._capture = { + _VGG19_PRE_RELU_INDICES[name]: weight for name, weight in weights.items() + } + + def train(self, mode: bool = True) -> "VGGPerceptualLoss": + """Keep the VGG features frozen in eval mode regardless of parent mode.""" + super().train(mode) + self.features.eval() + return self + + def _extract(self, x: Tensor) -> dict[int, Tensor]: + """Return the pre-ReLU feature maps at the configured indices.""" + x = (x - self.mean) / self.std + feats: dict[int, Tensor] = {} + for i, layer in enumerate(self.features): + x = layer(x) + if i in self._capture: + feats[i] = x + return feats + + def forward(self, fake: Tensor, real: Tensor) -> Tensor: + """Compute the weighted perceptual loss between ``fake`` and ``real``. + + Args: + fake: Generated RGB image ``[N, 3, H, W]`` in ``[0, 1]`` (keeps grad). + real: Target RGB image ``[N, 3, H, W]`` in ``[0, 1]``. + + Returns: + Scalar perceptual loss. + """ + param_dtype = self.mean.dtype + fake_feats = self._extract(fake.to(param_dtype)) + with torch.no_grad(): + real_feats = self._extract(real.to(param_dtype)) + loss = fake.new_zeros(()) + for index, weight in self._capture.items(): + if self.criterion == "l1": + term = F.l1_loss(fake_feats[index], real_feats[index]) + else: + term = F.mse_loss(fake_feats[index], real_feats[index]) + loss = loss + weight * term + return loss + + +@dataclass +class VGGPerceptualLossConfig: + """Configuration for :class:`VGGPerceptualLoss`.""" + + layer_weights: dict[str, float] = field( + default_factory=lambda: dict(_DEFAULT_LAYER_WEIGHTS) + ) + criterion: str = "l1" + + def build(self) -> VGGPerceptualLoss: + """Build the perceptual loss module.""" + return VGGPerceptualLoss( + layer_weights=self.layer_weights, + criterion=self.criterion, + ) diff --git a/olmoearth_pretrain/train/train_module/naip_gan.py b/olmoearth_pretrain/train/train_module/naip_gan.py new file mode 100644 index 000000000..46a9037f5 --- /dev/null +++ b/olmoearth_pretrain/train/train_module/naip_gan.py @@ -0,0 +1,903 @@ +"""Train module for the auxiliary NAIP conditional pix2pix-GAN objective. + +Extends :class:`LatentMIMTrainModule` with a generator (part of the model, +trained by the main optimizer alongside the encoder) and a conditional +discriminator (held here with its own optimizer). The discriminator loss only +updates the discriminator; the generator loss (adversarial + L1) backprops +through the encoder. +""" + +from dataclasses import dataclass, field +from logging import getLogger +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist +import torch.distributed.checkpoint.state_dict as dist_cp_sd +import torch.nn.functional as F +from olmo_core.distributed.parallel import DataParallelConfig +from olmo_core.distributed.utils import ( + get_local_rank, + get_local_tensor, + get_rank, + get_world_size, +) +from olmo_core.optim import OptimConfig +from olmo_core.optim.scheduler import Scheduler +from olmo_core.train.common import ReduceType + +from olmoearth_pretrain.config import Config +from olmoearth_pretrain.data.constants import MISSING_VALUE, Modality +from olmoearth_pretrain.data.normalize import load_computed_config +from olmoearth_pretrain.data.transform import TransformConfig +from olmoearth_pretrain.datatypes import ( + MaskedOlmoEarthSample, + MaskValue, + TokensAndMasks, +) +from olmoearth_pretrain.nn.naip_gan import ( + NaipGanModel, + discriminator_adversarial_loss, + generator_adversarial_loss, +) +from olmoearth_pretrain.nn.perceptual import VGGPerceptualLossConfig +from olmoearth_pretrain.nn.pooling import PoolingType, pool_unmasked_tokens +from olmoearth_pretrain.nn.utils import unpack_encoder_output +from olmoearth_pretrain.train.loss import LossConfig +from olmoearth_pretrain.train.masking import MaskingConfig +from olmoearth_pretrain.train.train_module.latent_mim import ( + LatentMIMTrainModule, + LatentMIMTrainModuleConfig, +) +from olmoearth_pretrain.train.utils import split_masked_batch + +logger = getLogger(__name__) + + +@dataclass +class NaipGanTrainModuleConfig(LatentMIMTrainModuleConfig): + """Configuration for :class:`NaipGanTrainModule`. + + Args: + discriminator_config: Config for the conditional discriminator. + disc_optim_config: Optimizer config for the discriminator. + lambda_adv: Weight on the generator adversarial loss. + lambda_l1: Weight on the generator L1 reconstruction loss. + lambda_perceptual: Weight on the generator VGG perceptual loss (0 + disables it and skips building the VGG network). Like the L1 term it + is active from step 0 and provides the mid-frequency content signal + that ESRGAN / Real-ESRGAN rely on. + perceptual_config: Config for the VGG perceptual loss. Defaults to the + ESRGAN layer weights when ``lambda_perceptual > 0``. + gan_loss_type: Adversarial loss variant ("hinge" or "bce"). + gan_warmup_steps: Number of steps before the adversarial terms turn on + (the L1 and perceptual reconstruction terms are active from step 0). + naip_modality: Name of the NAIP modality to predict. + image_log_interval: Log generated-vs-real NAIP images to W&B every this + many steps (0 disables image logging). + num_log_images: Max number of paired images to upload each time. + naip_denorm_std_multiplier: std multiplier used to invert the NAIP + normalization for display (must match the dataset Normalizer). + discriminator_cond_source: What the discriminator conditions on. One of + "online_pooled" (default; the online-encoder pooled embedding of the + masked input), "online_unmasked_pooled" (the online-encoder pooled + embedding from a full-depth forward on the unmasked input), or + "raw_sentinel2" (the raw Sentinel-2 image, conditioned on even when + it is masked for the online encoder). The discriminator's + ``cond_mode`` must match: "image" for "raw_sentinel2", "embedding" + otherwise. + cond_modality: Modality name used as the raw image condition when + ``discriminator_cond_source == "raw_sentinel2"``. The full temporal + stack of all its bands is fed to the discriminator. + """ + + discriminator_config: Config | None = None + disc_optim_config: OptimConfig | None = None + lambda_adv: float = 0.1 + lambda_l1: float = 10.0 + lambda_perceptual: float = 0.0 + perceptual_config: VGGPerceptualLossConfig | None = None + gan_loss_type: str = "hinge" + gan_warmup_steps: int = 0 + naip_modality: str = field(default_factory=lambda: Modality.NAIP_10.name) + image_log_interval: int = 1000 + num_log_images: int = 10 + naip_denorm_std_multiplier: float = 2.0 + discriminator_cond_source: str = "online_pooled" + cond_modality: str = field(default_factory=lambda: Modality.SENTINEL2_L2A.name) + + def build( + self, + model: NaipGanModel, + device: torch.device | None = None, + ) -> "NaipGanTrainModule": + """Build the corresponding :class:`NaipGanTrainModule`.""" + # The discriminator (and its optimizer) is only needed for the adversarial + # branch. When ``lambda_adv == 0`` the generator is trained by the + # reconstruction terms alone (pure L1 / perceptual), so it may be omitted. + if self.lambda_adv > 0: + if self.discriminator_config is None: + raise ValueError("discriminator_config is required when lambda_adv > 0") + if self.disc_optim_config is None: + raise ValueError("disc_optim_config is required when lambda_adv > 0") + kwargs = self.prepare_kwargs() + return NaipGanTrainModule( + model=model, + device=device, + **kwargs, + ) + + +class NaipGanTrainModule(LatentMIMTrainModule): + """LatentMIM train module with an auxiliary NAIP conditional GAN branch.""" + + def __init__( + self, + model: NaipGanModel, + optim_config: OptimConfig, + transform_config: TransformConfig, + masking_config: MaskingConfig, + loss_config: LossConfig, + rank_microbatch_size: int, + token_exit_cfg: dict[str, int], + discriminator_config: Config | None = None, + disc_optim_config: OptimConfig | None = None, + lambda_adv: float = 0.1, + lambda_l1: float = 10.0, + lambda_perceptual: float = 0.0, + perceptual_config: VGGPerceptualLossConfig | None = None, + gan_loss_type: str = "hinge", + gan_warmup_steps: int = 0, + naip_modality: str = Modality.NAIP_10.name, + image_log_interval: int = 1000, + num_log_images: int = 10, + naip_denorm_std_multiplier: float = 2.0, + discriminator_cond_source: str = "online_pooled", + cond_modality: str = Modality.SENTINEL2_L2A.name, + mae_loss_config: LossConfig | None = None, + compile_model: bool = False, + dp_config: DataParallelConfig | None = None, + compile_loss: bool = False, + autocast_precision: torch.dtype | None = None, + max_grad_norm: float | None = None, + scheduler: Scheduler | None = None, + device: torch.device | None = None, + state_dict_save_opts: dist_cp_sd.StateDictOptions | None = None, + state_dict_load_opts: dist_cp_sd.StateDictOptions | None = None, + ema_decay: tuple[float, float] = (0.996, 1.0), + regularizer_config: LossConfig | None = None, + find_unused_parameters: bool = True, + ): + """Initialize the train module (see config for arg descriptions).""" + super().__init__( + model=model, + optim_config=optim_config, + transform_config=transform_config, + masking_config=masking_config, + loss_config=loss_config, + rank_microbatch_size=rank_microbatch_size, + token_exit_cfg=token_exit_cfg, + mae_loss_config=mae_loss_config, + compile_model=compile_model, + dp_config=dp_config, + compile_loss=compile_loss, + autocast_precision=autocast_precision, + max_grad_norm=max_grad_norm, + scheduler=scheduler, + device=device, + state_dict_save_opts=state_dict_save_opts, + state_dict_load_opts=state_dict_load_opts, + ema_decay=ema_decay, + regularizer_config=regularizer_config, + find_unused_parameters=find_unused_parameters, + ) + self.lambda_adv = lambda_adv + self.lambda_l1 = lambda_l1 + self.lambda_perceptual = lambda_perceptual + self.gan_loss_type = gan_loss_type + self.gan_warmup_steps = gan_warmup_steps + self.naip_modality = naip_modality + self.image_log_interval = image_log_interval + self.num_log_images = num_log_images + self.discriminator_cond_source = discriminator_cond_source + self.cond_modality = cond_modality + valid_sources = ("online_pooled", "online_unmasked_pooled", "raw_sentinel2") + if discriminator_cond_source not in valid_sources: + raise ValueError( + f"discriminator_cond_source must be one of {valid_sources}, " + f"got {discriminator_cond_source}" + ) + + # Precompute per-band (R, G, B) min/max for inverse min-max denormalization + # of NAIP for W&B display, matching the dataset's computed normalization. + computed_config = load_computed_config()[self.naip_modality] + rgb_bands = Modality.get(self.naip_modality).band_order[:3] + means = np.array([computed_config[b]["mean"] for b in rgb_bands]) + stds = np.array([computed_config[b]["std"] for b in rgb_bands]) + self._naip_rgb_min = means - naip_denorm_std_multiplier * stds + self._naip_rgb_max = means + naip_denorm_std_multiplier * stds + # Torch versions on-device for the differentiable RGB[0,1] conversion the + # perceptual loss needs (the numpy versions above stay for W&B display). + self._naip_rgb_min_t = torch.tensor( + self._naip_rgb_min, dtype=torch.float32, device=self.device + ).reshape(1, 3, 1, 1) + self._naip_rgb_max_t = torch.tensor( + self._naip_rgb_max, dtype=torch.float32, device=self.device + ).reshape(1, 3, 1, 1) + + # The discriminator is intentionally NOT part of ``self.model`` so it is + # excluded from the main optimizer (its loss must not update the encoder). + # It is only built for the adversarial branch; when ``discriminator_config`` + # is None (pure L1 / perceptual training, ``lambda_adv == 0``) the module + # runs without any discriminator and all discriminator paths are skipped. + self.discriminator = None + self.disc_optimizer = None + if discriminator_config is not None: + self.discriminator = discriminator_config.build() + self.discriminator.to(self.device) + # The discriminator's condition input must match the chosen source: a + # raw image ("image") for raw Sentinel-2, otherwise a pooled embedding. + expected_cond_mode = ( + "image" if discriminator_cond_source == "raw_sentinel2" else "embedding" + ) + disc_cond_mode = getattr(self.discriminator, "cond_mode", "embedding") + if disc_cond_mode != expected_cond_mode: + raise ValueError( + f"discriminator_cond_source={discriminator_cond_source!r} requires " + f"the discriminator cond_mode={expected_cond_mode!r}, but got " + f"{disc_cond_mode!r}" + ) + # The discriminator is small and replicated (not sharded); rather than + # wrapping it in DDP we keep it as a plain module and all-reduce its + # gradients manually once per step (see + # ``_all_reduce_discriminator_grads``). Broadcast the initial weights + # from the first rank of the data-parallel group so every replica + # starts identical (this is what DDP would do at construction time). + if ( + self._dp_config is not None + and get_world_size(self.dp_process_group) > 1 + ): + src = dist.get_global_rank(self.dp_process_group, 0) + for tensor in [ + *self.discriminator.parameters(), + *self.discriminator.buffers(), + ]: + dist.broadcast( + tensor.detach(), src=src, group=self.dp_process_group + ) + assert disc_optim_config is not None + self.disc_optimizer = disc_optim_config.build(self.discriminator) + elif lambda_adv > 0: + raise ValueError("discriminator_config is required when lambda_adv > 0") + + # Perceptual loss network: a frozen VGG held here (like the discriminator, + # it is NOT part of ``self.model`` and has no optimizer). Its weights are + # deterministic pretrained ImageNet weights, so every rank is identical + # without any broadcast. Gradients flow through it into the generator. + self.perceptual_loss = None + if self.lambda_perceptual > 0: + cfg = perceptual_config or VGGPerceptualLossConfig() + self.perceptual_loss = cfg.build() + self.perceptual_loss.to(self.device) + self.perceptual_loss.eval() + for p in self.perceptual_loss.parameters(): + p.requires_grad_(False) + + self.total_loss_name = f"{self.total_loss_name}+G_l1" + if self.discriminator is not None: + self.total_loss_name = f"{self.total_loss_name}+G_adv" + if self.lambda_perceptual > 0: + self.total_loss_name = f"{self.total_loss_name}+G_perceptual" + + def zero_grads(self) -> None: + """Zero gradients for both the main optimizer and the discriminator.""" + super().zero_grads() + if self.disc_optimizer is not None: + self.disc_optimizer.zero_grad(set_to_none=True) + + def _all_reduce_discriminator_grads(self) -> None: + """Average the discriminator gradients across the data-parallel group. + + The discriminator is a plain (non-DDP) replicated module, so its + gradients are accumulated purely locally during backward. This performs + the single explicit all-reduce that keeps every replica's update (and + therefore its weights) identical. Ranks that produced no discriminator + gradient this step (e.g. no valid NAIP) contribute zeros so that every + rank still participates in the collective; the sum is divided by the full + data-parallel world size, matching standard DDP averaging. + """ + if self.discriminator is None: + return + if self._dp_config is None: + return + world_size = get_world_size(self.dp_process_group) + if world_size <= 1: + return + for p in self.discriminator.parameters(): + if p.grad is None: + p.grad = torch.zeros_like(p) + dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, group=self.dp_process_group) + p.grad /= world_size + + def _sync_discriminator_buffers(self) -> None: + """Broadcast the discriminator's buffers from data-parallel rank 0. + + The discriminator is a plain replicated module whose gradients are + all-reduced (so its parameters stay identical across ranks), but any + buffers updated in forward drift independently per rank. In particular, + spectral normalization's power-iteration vectors are updated a different + number of times on ranks that skip the discriminator forward this step + (e.g. no valid NAIP), so the effective normalization diverges. This + rebroadcast keeps every replica bit-identical, matching DDP's default + ``broadcast_buffers=True`` behavior (a no-op when there are no buffers). + """ + if self.discriminator is None: + return + if self._dp_config is None: + return + if get_world_size(self.dp_process_group) <= 1: + return + src = dist.get_global_rank(self.dp_process_group, 0) + for buf in self.discriminator.buffers(): + dist.broadcast(buf.detach(), src=src, group=self.dp_process_group) + + def model_forward( # type: ignore[override] + self, + batch: MaskedOlmoEarthSample, + patch_size: int, + token_exit_cfg: dict[str, int], + ) -> tuple[ + torch.Tensor, + TokensAndMasks, + TokensAndMasks, + TokensAndMasks, + torch.Tensor, + torch.Tensor, + ]: + """Run a forward pass returning the SSL loss plus the GAN tensors.""" + with self._model_forward_context(): + ( + latent, + decoded, + _, + reconstructed, + extra_metrics, + pooled, + fake_naip, + ) = self.model(batch, patch_size) + if extra_metrics is not None: + self.log_extra_metrics(extra_metrics) + with torch.no_grad(): + logger.info("Target Encoder forward pass...") + output_dict = self.model.target_encoder.forward( + batch.unmask(), + patch_size=patch_size, + token_exit_cfg=token_exit_cfg, + ) + target_output, _, _ = unpack_encoder_output(output_dict) + loss = self.loss_fn(decoded, target_output) + if self.mae_loss is not None: + loss += self.mae_loss.compute(reconstructed, batch) + return loss, latent, decoded, target_output, pooled, fake_naip + + def _extract_real_naip( + self, batch: MaskedOlmoEarthSample + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Extract the real NAIP image and a per-instance validity mask. + + Returns: + (image ``[B, C, Hn, Wn]``, valid ``[B]`` bool) or (None, None) if the + NAIP modality is absent from the batch. + """ + naip = getattr(batch, self.naip_modality, None) + mask = getattr( + batch, + MaskedOlmoEarthSample.get_masked_modality_name(self.naip_modality), + None, + ) + if naip is None or mask is None: + return None, None + batch_size = naip.shape[0] + valid = (mask != MaskValue.MISSING.value).reshape(batch_size, -1).any(dim=1) + # NAIP is not multitemporal (T == 1); drop time and move channels first. + naip_img = naip[:, :, :, 0, :].permute(0, 3, 1, 2).contiguous() + return naip_img, valid + + def _naip_to_rgb01(self, x: torch.Tensor) -> np.ndarray: + """Invert NAIP normalization and return displayable RGB in ``[0, 1]``. + + Args: + x: Normalized NAIP tensor of shape ``[N, C, H, W]`` (C >= 3). + + Returns: + A ``[N, H, W, 3]`` float array in ``[0, 1]``. + """ + rgb = x[:, :3].detach().float().cpu().numpy() # [N, 3, H, W] + min_vals = self._naip_rgb_min.reshape(1, 3, 1, 1) + max_vals = self._naip_rgb_max.reshape(1, 3, 1, 1) + raw = rgb * (max_vals - min_vals) + min_vals + rgb01 = np.clip(raw / 255.0, 0.0, 1.0) + return np.transpose(rgb01, (0, 2, 3, 1)) # [N, H, W, 3] + + def _naip_to_rgb01_tensor(self, x: torch.Tensor) -> torch.Tensor: + """Differentiably invert NAIP normalization to RGB ``[0, 1]`` for VGG. + + Unlike :meth:`_naip_to_rgb01` (numpy, for display) this keeps the graph + so the perceptual loss can backprop into the generator. The raw pixel + range is mapped linearly to ``[0, 1]`` and softly clamped with + ``clamp`` only at the boundaries (values already inside the range keep a + unit gradient). + + Args: + x: Normalized NAIP tensor ``[N, C, H, W]`` (uses the first 3 bands). + + Returns: + RGB tensor ``[N, 3, H, W]`` in ``[0, 1]``. + """ + rgb = x[:, :3].float() + min_vals = self._naip_rgb_min_t.to(rgb.device) + max_vals = self._naip_rgb_max_t.to(rgb.device) + raw = rgb * (max_vals - min_vals) + min_vals + return (raw / 255.0).clamp(0.0, 1.0) + + def _maybe_log_naip_images( + self, batch: MaskedOlmoEarthSample, fake_naip: torch.Tensor | None + ) -> None: + """Occasionally upload paired real/generated NAIP images to W&B.""" + if self.image_log_interval <= 0 or fake_naip is None: + return + if self.trainer.global_step % self.image_log_interval != 0: + return + if get_rank() != 0: + return + + real_img, valid = self._extract_real_naip(batch) + if real_img is None or valid is None or not bool(valid.any()): + return + + wandb = self._get_wandb() + if wandb is None: + return + + num = min(self.num_log_images, int(valid.sum().item())) + fake_v = fake_naip[valid][:num] + real_v = F.interpolate( + real_img[valid][:num].to(fake_v.dtype), + size=fake_v.shape[-2:], + mode="bilinear", + align_corners=False, + ) + fake_rgb = self._naip_to_rgb01(fake_v) + real_rgb = self._naip_to_rgb01(real_v) + + images = [] + for i in range(num): + # Real on the left, generated on the right. + pair = np.concatenate([real_rgb[i], fake_rgb[i]], axis=1) + images.append(wandb.Image(pair, caption="left: real | right: generated")) + wandb.log( + {"gan/naip_real_vs_fake": images}, + step=self.trainer.global_step, + ) + + def _get_wandb(self) -> Any | None: + """Return the wandb module from the OlmoEarth W&B callback, if enabled.""" + from olmoearth_pretrain.train.callbacks.wandb import OlmoEarthWandBCallback + + for callback in self.trainer._iter_callbacks(): + if isinstance(callback, OlmoEarthWandBCallback) and callback.enabled: + return callback.wandb + return None + + def train_batch( + self, + batch: tuple[int, MaskedOlmoEarthSample], + dry_run: bool = False, + ) -> None: + """Train a batch with the SSL objective plus the NAIP GAN branch.""" + if not dry_run: + self.update_target_encoder() + self.model.train() + if self.discriminator is not None: + self.discriminator.train() + + total_batch_loss = torch.zeros([], device=self.device) + total_batch_reg = torch.zeros([], device=self.device) + total_d_loss = torch.zeros([], device=self.device) + total_g_adv = torch.zeros([], device=self.device) + total_g_l1 = torch.zeros([], device=self.device) + total_g_perceptual = torch.zeros([], device=self.device) + + patch_size = batch[0] + batch_data = batch[1] + + adversarial_active = ( + self.discriminator is not None + and not dry_run + and self.trainer.global_step >= self.gan_warmup_steps + ) + + masked_microbatches = split_masked_batch(batch_data, self.rank_microbatch_size) + num_microbatches = len(masked_microbatches) + + for microbatch_idx in range(num_microbatches): + with self._train_microbatch_context(microbatch_idx, num_microbatches): + microbatch_masked = masked_microbatches[microbatch_idx] + logger.info( + f"Training microbatch {microbatch_idx} of {num_microbatches} " + f"with batch size {microbatch_masked.batch_size} on rank {get_local_rank()}" + ) + masked_batch = microbatch_masked.to_device(self.device) + + loss, latent, decoded, target_output, pooled, fake_naip = ( + self.model_forward(masked_batch, patch_size, self.token_exit_cfg) + ) + + if not dry_run and microbatch_idx == 0: + with torch.no_grad(): + self._maybe_log_naip_images(masked_batch, fake_naip) + + reg_term = self.compute_regularization(latent) + if reg_term is not None: + loss = loss + reg_term + total_batch_reg += ( + get_local_tensor(reg_term.detach()) / num_microbatches + ) + + # Skip the condition extraction entirely without a discriminator + # (pure L1 / perceptual training): it is only consumed by the + # discriminator/adversarial paths, and the ``online_unmasked_pooled`` + # source would otherwise run an extra full-depth encoder forward. + if self.discriminator is None: + d_cond, d_cond_valid, d_cond_time_mask = None, None, None + else: + d_cond, d_cond_valid, d_cond_time_mask = ( + self._extract_discriminator_cond( + masked_batch, pooled, patch_size + ) + ) + + d_loss = self._maybe_discriminator_loss( + masked_batch, + d_cond, + d_cond_valid, + fake_naip, + num_microbatches, + adversarial_active, + patch_size, + d_cond_time_mask, + ) + gen_loss, g_l1_val, g_adv_val, g_perceptual_val = self._generator_loss( + masked_batch, + d_cond, + d_cond_valid, + fake_naip, + adversarial_active, + patch_size, + d_cond_time_mask, + ) + total_g_l1 += g_l1_val / num_microbatches + total_g_adv += g_adv_val / num_microbatches + total_g_perceptual += g_perceptual_val / num_microbatches + + # Anchor the generator to the loss with zero weight so its + # parameters always receive a gradient on every rank. When a rank + # has no valid NAIP, the generator is otherwise disconnected from + # the loss (both the L1 and adversarial terms are skipped), so its + # gradient flow -- and therefore the root model's FSDP gradient + # reduce-scatter -- would be asymmetric across ranks and could + # deadlock the collective. The term is exactly zero and does not + # change any real gradient. + if fake_naip is not None: + loss = loss + 0.0 * fake_naip.sum() + + loss = (loss + gen_loss) / num_microbatches + total_batch_loss += get_local_tensor(loss.detach()) + + if torch.isnan(loss).any() or torch.isinf(loss).any(): + logger.warning( + f"NaN or Inf detected in loss at microbatch {microbatch_idx}, " + f"stopping training for this batch." + ) + print(f"rank {get_local_rank()} has nan or inf") + + if d_loss is not None: + total_d_loss += get_local_tensor(d_loss.detach()) + d_loss.backward() + loss.backward() + + if not dry_run and adversarial_active: + # The discriminator accumulated gradients locally across microbatches; + # average them across the data-parallel group once, then step. Gating + # on ``adversarial_active`` (which is derived from the global step and + # so is identical on every rank) guarantees all ranks participate in + # the collective together, even if some had no valid NAIP this step. + self._all_reduce_discriminator_grads() + assert self.disc_optimizer is not None + self.disc_optimizer.step() + # Keep the replicated discriminator bit-identical across ranks: its + # buffers (e.g. spectral-norm power-iteration vectors) can drift when + # ranks run the discriminator forward a different number of times. + self._sync_discriminator_buffers() + + if dry_run: + if self.disc_optimizer is not None: + self.disc_optimizer.zero_grad(set_to_none=True) + return + + self.trainer.record_metric( + f"train/{self.total_loss_name}", total_batch_loss, ReduceType.mean + ) + self.trainer.record_metric("train/G_l1", total_g_l1, ReduceType.mean) + if self.discriminator is not None: + self.trainer.record_metric("train/G_adv", total_g_adv, ReduceType.mean) + self.trainer.record_metric("train/D_loss", total_d_loss, ReduceType.mean) + if self.perceptual_loss is not None: + self.trainer.record_metric( + "train/G_perceptual", total_g_perceptual, ReduceType.mean + ) + self.log_regularization(total_batch_reg) + + del batch, batch_data + del masked_batch + del latent, decoded, target_output, pooled, fake_naip + + def _extract_discriminator_cond( + self, + batch: MaskedOlmoEarthSample, + pooled: torch.Tensor, + patch_size: int, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Return the discriminator condition, validity mask, and time mask. + + The condition depends on ``discriminator_cond_source``: + + * ``online_pooled``: the online-encoder pooled embedding of the masked + input ``[B, H, W, D]``. + * ``online_unmasked_pooled``: the online-encoder pooled embedding from a + full-depth forward on the unmasked input ``[B, H, W, D]`` (run under + ``no_grad``; an extra encoder forward per step). + * ``raw_sentinel2``: the raw Sentinel-2 temporal stack + ``[B, T, C, Hs, Ws]``. + + The condition is always detached where it feeds the discriminator, so the + encoder never receives gradient through the conditioning path (only the + generator's L1/adversarial loss through the generated image trains it). + + Returns: + ``(cond, cond_valid, cond_time_mask)``. ``cond`` is None when the raw + condition is entirely absent from the batch; ``cond_valid`` is a + ``[B]`` bool mask (None means every instance is valid); + ``cond_time_mask`` is a ``[B, T]`` bool mask of valid timesteps for + the raw stack (None for the pooled embedding modes). + """ + source = self.discriminator_cond_source + if source == "online_pooled": + return pooled, None, None + if source == "online_unmasked_pooled": + # Full-depth online-encoder forward on the unmasked input. The + # condition is detached at the discriminator, so this runs under + # no_grad (forward only). + with self._model_forward_context(), torch.no_grad(): + output_dict = self.model.encoder.forward( + batch.unmask(), + patch_size=patch_size, + token_exit_cfg=None, + ) + online_unmasked_output, _, _ = unpack_encoder_output(output_dict) + pooled_unmasked = pool_unmasked_tokens( + online_unmasked_output, PoolingType.MEAN, spatial_pooling=True + ) + return pooled_unmasked, None, None + # raw_sentinel2 + return self._extract_raw_s2_cond(batch) + + def _extract_raw_s2_cond( + self, batch: MaskedOlmoEarthSample + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Extract the raw image temporal stack and validity masks. + + The raw modality tensor retains its full pixel values regardless of the + online-encoder mask, so an instance can be conditioned on even when the + modality is masked for the online encoder (any non-MISSING mask counts as + valid). All bands and all timesteps are returned; missing entries (the + MISSING_VALUE sentinel) are zeroed and excluded from the temporal mean via + the returned per-timestep validity mask. + + Returns: + ``(stack [B, T, C, H, W], valid [B], time_valid [B, T])`` or + ``(None, None, None)`` if the modality is absent from the batch. + """ + modality = self.cond_modality + raw = getattr(batch, modality, None) + mask = getattr( + batch, + MaskedOlmoEarthSample.get_masked_modality_name(modality), + None, + ) + if raw is None or mask is None: + return None, None, None + batch_size = raw.shape[0] + # [B, H, W, T, C] -> [B, T, C, H, W] (all bands). + stack = raw.permute(0, 3, 4, 1, 2) + present = stack != MISSING_VALUE + # Zero the sentinel so the per-timestep convs don't see huge values. + stack = torch.where(present, stack, torch.zeros_like(stack)) + # Per-timestep validity (any present pixel/band). + time_valid = present.reshape(batch_size, stack.shape[1], -1).any(dim=2) + # Per-instance validity. + valid = (mask != MaskValue.MISSING.value).reshape(batch_size, -1).any(dim=1) + return stack, valid, time_valid + + def _maybe_discriminator_loss( + self, + batch: MaskedOlmoEarthSample, + cond: torch.Tensor | None, + cond_valid: torch.Tensor | None, + fake_naip: torch.Tensor, + num_microbatches: int, + adversarial_active: bool, + patch_size: int, + cond_time_mask: torch.Tensor | None = None, + ) -> torch.Tensor | None: + """Compute the (scaled) discriminator loss on detached inputs. + + Returns None when the adversarial phase is inactive, NAIP is absent, or no + instance has both valid NAIP and a valid condition. Only the + discriminator receives gradients from this loss. + """ + if not adversarial_active or fake_naip is None or cond is None: + return None + # ``adversarial_active`` implies a discriminator exists. + assert self.discriminator is not None + real_img, valid = self._extract_real_naip(batch) + if real_img is None or valid is None: + return None + if cond_valid is not None: + valid = valid & cond_valid + if not bool(valid.any()): + return None + with self._model_forward_context(): + fake_v = fake_naip[valid].detach() + cond_v = cond[valid].detach() + time_v = None if cond_time_mask is None else cond_time_mask[valid] + real_v = F.interpolate( + real_img[valid].to(fake_naip.dtype), + size=fake_naip.shape[-2:], + mode="bilinear", + align_corners=False, + ) + # Single discriminator forward on the concatenated real+fake batch + # (one forward is slightly cheaper than two separate passes). + images = torch.cat([real_v, fake_v], dim=0) + conds = torch.cat([cond_v, cond_v], dim=0) + time_mask = None if time_v is None else torch.cat([time_v, time_v], dim=0) + logits = self.discriminator( + images, + conds, + patch_size=patch_size, + cond_time_mask=time_mask, + ) + real_logits, fake_logits = logits.chunk(2, dim=0) + d_loss = discriminator_adversarial_loss( + real_logits, fake_logits, self.gan_loss_type + ) + return d_loss / num_microbatches + + def _generator_loss( + self, + batch: MaskedOlmoEarthSample, + cond: torch.Tensor | None, + cond_valid: torch.Tensor | None, + fake_naip: torch.Tensor, + adversarial_active: bool, + patch_size: int, + cond_time_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute the generator loss (L1 + perceptual always, adversarial once warmed up). + + Gradients from this loss flow into the generator and encoder but not the + discriminator (its parameters are frozen during the forward pass here). + + Returns: + (generator_loss, g_l1_value, g_adv_value, g_perceptual_value); the + loss is zero when NAIP is absent from the batch. + """ + zero = torch.zeros([], device=self.device) + if fake_naip is None: + return zero, zero, zero, zero + real_img, valid = self._extract_real_naip(batch) + if real_img is None or valid is None: + return zero, zero, zero, zero + if cond_valid is not None: + valid = valid & cond_valid + if not bool(valid.any()): + return zero, zero, zero, zero + + with self._model_forward_context(): + fake_v = fake_naip[valid] + real_v = F.interpolate( + real_img[valid].to(fake_naip.dtype), + size=fake_naip.shape[-2:], + mode="bilinear", + align_corners=False, + ) + g_l1 = F.l1_loss(fake_v, real_v) + gen_loss = self.lambda_l1 * g_l1 + g_perceptual = zero + if self.perceptual_loss is not None: + # Compare RGB in [0, 1] (ImageNet-normalized inside the VGG). + g_perceptual = self.perceptual_loss( + self._naip_to_rgb01_tensor(fake_v), + self._naip_to_rgb01_tensor(real_v), + ) + gen_loss = gen_loss + self.lambda_perceptual * g_perceptual + g_adv = zero + if adversarial_active and cond is not None: + # ``adversarial_active`` implies a discriminator exists. + assert self.discriminator is not None + cond_v = cond[valid] + time_v = None if cond_time_mask is None else cond_time_mask[valid] + # Freeze the discriminator so the generator loss doesn't update + # it: without this, ``loss.backward()`` (which includes g_adv) + # would accumulate adversarial gradients into the discriminator's + # params, which would then be applied by ``disc_optimizer``. + for p in self.discriminator.parameters(): + p.requires_grad_(False) + # Detach the condition so the adversarial gradient only reaches + # the encoder through the generated image (fake_v), not through + # the conditioning tokens the discriminator ingests. This stops + # the encoder from "cheating" the discriminator by reshaping the + # condition embedding instead of improving the generated image. + fake_logits = self.discriminator( + fake_v, + cond_v.detach(), + patch_size=patch_size, + cond_time_mask=time_v, + ) + for p in self.discriminator.parameters(): + p.requires_grad_(True) + g_adv = generator_adversarial_loss(fake_logits, self.gan_loss_type) + gen_loss = gen_loss + self.lambda_adv * g_adv + return ( + gen_loss, + get_local_tensor(g_l1.detach()), + get_local_tensor(g_adv.detach()), + get_local_tensor(g_perceptual.detach()), + ) + + def _get_state_dict( + self, sd_options: dist_cp_sd.StateDictOptions + ) -> dict[str, Any]: + """Include the discriminator + its optimizer for clean resume.""" + sd = super()._get_state_dict(sd_options) + if self.discriminator is not None: + assert self.disc_optimizer is not None + sd["discriminator"] = dist_cp_sd.get_model_state_dict( + self.discriminator, options=sd_options + ) + sd["disc_optim"] = dist_cp_sd.get_optimizer_state_dict( + self.discriminator, self.disc_optimizer, options=sd_options + ) + return sd + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load model + optim, then the discriminator state if present.""" + super().load_state_dict(state_dict) + if self.discriminator is None: + return + assert self.disc_optimizer is not None + if "discriminator" in state_dict: + dist_cp_sd.set_model_state_dict( + self.discriminator, + state_dict["discriminator"], + options=self.state_dict_load_opts, + ) + if "disc_optim" in state_dict: + dist_cp_sd.set_optimizer_state_dict( + self.discriminator, + self.disc_optimizer, + state_dict["disc_optim"], + options=self.state_dict_load_opts, + ) diff --git a/scripts/official/v1_2/naip_gan.py b/scripts/official/v1_2/naip_gan.py new file mode 100644 index 000000000..d8453ad88 --- /dev/null +++ b/scripts/official/v1_2/naip_gan.py @@ -0,0 +1,240 @@ +"""v1.2 encoder + auxiliary NAIP conditional pix2pix-GAN branch. + +This is a single-view variant of the v1.2 baseline (``base.py``): it keeps the +latent-MIM patch-discrimination objective but adds a NAIP generator on top of +the encoder's pooled spatial embedding plus a conditional discriminator. NAIP +(``naip_10``) is a decode-only modality, so the generator must synthesize it +from the (masked) Sentinel-2 / other encode tokens. The discriminator here +conditions on the online encoder's pooled embedding from a full-depth forward on +the unmasked input (``online_unmasked_pooled``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan.py dry_run naip_gan local +""" + +import logging + +import base as v1_2_base +from olmo_core.config import DType +from olmo_core.distributed.parallel.data_parallel import ( + DataParallelConfig, + DataParallelType, +) +from olmo_core.optim import AdamWConfig +from olmo_core.optim.scheduler import CosWithWarmup +from olmo_core.train.config import TrainerConfig + +from olmoearth_pretrain.data.constants import Modality +from olmoearth_pretrain.data.dataloader import OlmoEarthDataLoaderConfig +from olmoearth_pretrain.data.dataset import OlmoEarthDatasetConfig +from olmoearth_pretrain.internal.experiment import ( + CommonComponents, + OlmoEarthVisualizeConfig, + SubCmd, + main, +) +from olmoearth_pretrain.internal.utils import MODEL_SIZE_ARGS +from olmoearth_pretrain.nn.naip_gan import ( + NaipDiscriminatorConfig, + NaipGanModelConfig, + NaipGeneratorConfig, +) +from olmoearth_pretrain.train.loss import LossConfig +from olmoearth_pretrain.train.masking import MaskingConfig +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +# The encoder patch size varies per batch (flexi-ViT, 1-8 here). The generator's +# learned unpatchify factor and the discriminator's condition unpatchify factor +# are a fixed "canonical" patch size (40 m/px tokens): each batch's token grid is +# bilinearly resampled to this canonical grid before the learned unpatchify, so +# any encoder patch size lands at the native NAIP resolution. Encoder patch sizes +# below the canonical are downsampled to it (lossy); those >= it are upsampled. +CANONICAL_PATCH_SIZE = 4 +MIN_PATCH_SIZE = 1 +MAX_PATCH_SIZE = 5 + +# NAIP is added as a decode-only modality on top of the v1.2 decode targets. +ONLY_DECODE_MODALITIES = [*v1_2_base.ONLY_DECODE_MODALITIES, Modality.NAIP_10.name] + +# Conv-trunk upsampling after the learned unpatchify. Set to the NAIP tile size +# factor so the generator output lands at native NAIP (2.5 m/px) resolution. +NAIP_UPSAMPLE_FACTOR = Modality.NAIP_10.image_tile_size_factor +# Generator per-stage channel widths (base 10 m/px, then 5 m/px, then 2.5 m/px): +# widened at the fine (post-upsampling) stages so more capacity lands where NAIP +# texture is synthesized. Upsampling between stages is sub-pixel (PixelShuffle). +GENERATOR_HIDDEN_SIZES = [256, 192, 192] +# Discriminator NAIP image stack: a stride-1 stem then two strided convs, each +# followed by two stride-1 refinement convs; image and condition features are +# fused at DISCRIMINATOR_FEATURE_CHANNELS by a head with two residual blocks. +DISCRIMINATOR_IMAGE_CHANNELS = [128, 256, 256] +DISCRIMINATOR_FEATURE_CHANNELS = 256 +DISCRIMINATOR_NUM_CONVS_PER_RESOLUTION = 2 +DISCRIMINATOR_NUM_HEAD_RES_BLOCKS = 2 +# The discriminator conditions on the online-encoder pooled embedding from a +# full-depth forward on the unmasked input. The embedding tokens are resampled to +# the canonical unpatchify factor (40 m/px) then a learned unpatchify expands them +# to the 10 m/px fusion grid; two convs (128 -> 256 -> 128) refine the condition +# after the unpatchify. +DISCRIMINATOR_COND_SOURCE = "online_unmasked_pooled" +DISCRIMINATOR_COND_UNPATCHIFY_FACTOR = CANONICAL_PATCH_SIZE +DISCRIMINATOR_COND_EMBEDDING_CHANNELS = [256] +DISCRIMINATOR_USE_PROJECTION = False +LAMBDA_ADV = 0.1 +LAMBDA_L1 = 10.0 +GAN_WARMUP_STEPS = 8000 + + +def _masking_config(tokenization_config=None) -> MaskingConfig: + return MaskingConfig( + strategy_config={ + "type": "random_time_with_decode", + "encode_ratio": 0.5, + "decode_ratio": 0.5, + "random_ratio": 0.5, + "only_decode_modalities": ONLY_DECODE_MODALITIES, + }, + tokenization_config=tokenization_config, + ) + + +def build_common_components( + script: str, cmd: SubCmd, run_name: str, cluster: str, overrides: list[str] +) -> CommonComponents: + """Build common components: v1.2 modalities plus NAIP.""" + config = v1_2_base.build_common_components( + script, cmd, run_name, cluster, overrides + ) + config.training_modalities = [ + *config.training_modalities, + Modality.NAIP_10.name, + ] + return config + + +def build_model_config(common: CommonComponents) -> NaipGanModelConfig: + """Build the NAIP GAN model config (v1.2 encoder/decoder + generator).""" + base_model = v1_2_base.build_size_model_config( + common, "base_shallow_decoder", v1_2_base.PATCH_EMBED_HIDDEN_SIZES + ) + # Flexi encoder: the patch size varies per batch within [MIN, MAX]. + base_model.encoder_config.min_patch_size = MIN_PATCH_SIZE + base_model.encoder_config.max_patch_size = MAX_PATCH_SIZE + embedding_size = ( + base_model.encoder_config.output_embedding_size + or base_model.encoder_config.embedding_size + ) + generator_config = NaipGeneratorConfig( + embedding_size=embedding_size, + patch_size=CANONICAL_PATCH_SIZE, + hidden_sizes=GENERATOR_HIDDEN_SIZES, + out_channels=Modality.NAIP_10.num_bands, + upsample_factor=NAIP_UPSAMPLE_FACTOR, + ) + return NaipGanModelConfig( + encoder_config=base_model.encoder_config, + decoder_config=base_model.decoder_config, + generator_config=generator_config, + ) + + +def build_train_module_config( + common: CommonComponents, +) -> NaipGanTrainModuleConfig: + """Build the NAIP GAN train module config.""" + model_size = MODEL_SIZE_ARGS["base_shallow_decoder"] + embedding_size = model_size["encoder_embedding_size"] + return NaipGanTrainModuleConfig( + optim_config=AdamWConfig(lr=0.0001, weight_decay=0.02, fused=False), + rank_microbatch_size=64, + masking_config=_masking_config(common.tokenization_config), + loss_config=LossConfig( + loss_config={ + "type": "modality_patch_discrimination_masked_negatives_vec", + "tau": 0.1, + "same_target_threshold": 0.999, + "mask_negatives_for_modalities": ONLY_DECODE_MODALITIES, + } + ), + discriminator_config=NaipDiscriminatorConfig( + embedding_size=embedding_size, + in_channels=Modality.NAIP_10.num_bands, + image_strided_conv_channels=DISCRIMINATOR_IMAGE_CHANNELS, + feature_channels=DISCRIMINATOR_FEATURE_CHANNELS, + num_convs_per_resolution=DISCRIMINATOR_NUM_CONVS_PER_RESOLUTION, + num_head_res_blocks=DISCRIMINATOR_NUM_HEAD_RES_BLOCKS, + use_projection=DISCRIMINATOR_USE_PROJECTION, + cond_unpatchify_factor=DISCRIMINATOR_COND_UNPATCHIFY_FACTOR, + cond_embedding_channels=DISCRIMINATOR_COND_EMBEDDING_CHANNELS, + ), + discriminator_cond_source=DISCRIMINATOR_COND_SOURCE, + disc_optim_config=AdamWConfig( + lr=0.0002, betas=(0.5, 0.999), weight_decay=0.0, fused=False + ), + lambda_adv=LAMBDA_ADV, + lambda_l1=LAMBDA_L1, + gan_warmup_steps=GAN_WARMUP_STEPS, + token_exit_cfg={modality: 0 for modality in common.training_modalities}, + max_grad_norm=1.0, + scheduler=CosWithWarmup(warmup_steps=8000), + ema_decay=(1.0, 1.0), + dp_config=DataParallelConfig( + name=DataParallelType.fsdp, + param_dtype=DType.bfloat16, + reduce_dtype=DType.float32, + ), + ) + + +def build_dataloader_config(common: CommonComponents) -> OlmoEarthDataLoaderConfig: + """Build the dataloader config (single masked view for latent-MIM).""" + return OlmoEarthDataLoaderConfig( + num_workers=16, + global_batch_size=512, + token_budget=2250, + prefetch_factor=4, + sampled_hw_p_list=list(range(1, 13)), + min_patch_size=MIN_PATCH_SIZE, + max_patch_size=MAX_PATCH_SIZE, + work_dir=common.save_folder, + seed=3622, + num_masked_views=1, + masking_config=_masking_config(common.tokenization_config), + ) + + +def build_dataset_config(common: CommonComponents) -> OlmoEarthDatasetConfig: + """Build the dataset config (a NAIP-containing h5py dataset).""" + return OlmoEarthDatasetConfig( + h5py_dir="/weka/dfive-default/helios/dataset/osm_sampling/h5py_data_w_missing_timesteps_zstd_3_128_x_4/cdl_landsat_naip_10_openstreetmap_raster_sentinel1_sentinel2_l2a_srtm_worldcereal_worldcover_wri_canopy_height_map/1138828", + training_modalities=common.training_modalities, + ) + + +WANDB_PROJECT = "2026_07_01_naip_gan" + + +def build_trainer_config(common: CommonComponents) -> TrainerConfig: + """Reuse the v1.2 trainer config, overriding the W&B project.""" + trainer_config = v1_2_base.build_trainer_config(common) + trainer_config.callbacks["wandb"].project = WANDB_PROJECT + return trainer_config + + +def build_visualize_config(common: CommonComponents) -> OlmoEarthVisualizeConfig: + """Reuse the v1.2 visualize config.""" + return v1_2_base.build_visualize_config(common) + + +if __name__ == "__main__": + main( + common_components_builder=build_common_components, + model_config_builder=build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=build_dataset_config, + dataloader_config_builder=build_dataloader_config, + trainer_config_builder=build_trainer_config, + visualize_config_builder=build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2.py b/scripts/official/v1_2/naip_gan_s2.py new file mode 100644 index 000000000..afa2a2bed --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2.py @@ -0,0 +1,53 @@ +"""v1.2 NAIP conditional GAN conditioned on the raw Sentinel-2 input. + +Instead of a pooled embedding, the discriminator conditions on the raw +Sentinel-2 temporal stack (``raw_sentinel2``): the full time series of all S2 +bands is embedded by non-strided per-timestep convs (3 layers), averaged over +the valid timesteps, then refined by more non-strided convs before fusing with +the NAIP image features. The generator and NAIP image path match +``naip_gan_unmasked_embed.py``; the projection term is off. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2.py dry_run naip_gan_s2 local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.data.constants import Modality +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +# Non-strided convs applied to each Sentinel-2 timestep before the mean over +# time (input bands -> 256 -> 256 -> 256 -> 256), then three more (256) after +# the mean pooling. +DISCRIMINATOR_COND_PRE_POOL_CHANNELS = [256, 256, 256, 256] +DISCRIMINATOR_COND_POST_POOL_CHANNELS = [256, 256, 256] + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the train module config conditioning on the raw Sentinel-2 stack.""" + config = naip_gan_base.build_train_module_config(common) + disc = config.discriminator_config + disc.cond_mode = "image" + disc.cond_in_channels = Modality.SENTINEL2_L2A.num_bands + disc.cond_image_pre_pool_channels = DISCRIMINATOR_COND_PRE_POOL_CHANNELS + disc.cond_image_post_pool_channels = DISCRIMINATOR_COND_POST_POOL_CHANNELS + config.discriminator_cond_source = "raw_sentinel2" + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_moregan.py b/scripts/official/v1_2/naip_gan_s2_moregan.py new file mode 100644 index 000000000..422e4343e --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_moregan.py @@ -0,0 +1,43 @@ +"""v1.2 NAIP raw-Sentinel-2 conditional GAN with a stronger adversarial term. + +Same as ``naip_gan_s2.py`` (raw Sentinel-2 condition) but rebalances the losses +toward the GAN objective: the L1 reconstruction weight is lowered +(``lambda_l1=2.0``) and the adversarial weight is raised (``lambda_adv=0.5``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_moregan.py dry_run naip_gan_s2_moregan local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2 as naip_gan_s2_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +LAMBDA_L1 = 2.0 +LAMBDA_ADV = 0.5 + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the raw-Sentinel-2 train module config with a stronger GAN term.""" + config = naip_gan_s2_base.build_train_module_config(common) + config.lambda_l1 = LAMBDA_L1 + config.lambda_adv = LAMBDA_ADV + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_percep_sn_bce.py b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce.py new file mode 100644 index 000000000..4ca5ecbf5 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce.py @@ -0,0 +1,50 @@ +"""v1.2 NAIP raw-Sentinel-2 GAN, ESRGAN-faithful recipe. + +Reference variant of the perceptual + spectral-norm + BCE sweep: raw +Sentinel-2 condition (no projection term), a spectrally-normalized +discriminator, a vanilla BCE adversarial loss, and the ESRGAN loss balance +(``lambda_l1=1.0``, ``lambda_perceptual=1.0``, ``lambda_adv=0.1``). This is the +closest match to the satlas super-resolution / Real-ESRGAN recipe. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_percep_sn_bce.py dry_run naip_gan_s2_percep_sn_bce local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2 as naip_gan_s2_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +# ESRGAN / satlas loss balance. +LAMBDA_L1 = 1.0 +LAMBDA_PERCEPTUAL = 1.0 +LAMBDA_ADV = 0.1 + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the raw-S2 config with perceptual loss, spectral norm, and BCE.""" + config = naip_gan_s2_base.build_train_module_config(common) + config.discriminator_config.use_spectral_norm = True + config.gan_loss_type = "bce" + config.lambda_l1 = LAMBDA_L1 + config.lambda_perceptual = LAMBDA_PERCEPTUAL + config.lambda_adv = LAMBDA_ADV + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_moregan.py b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_moregan.py new file mode 100644 index 000000000..740234321 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_moregan.py @@ -0,0 +1,46 @@ +"""v1.2 NAIP raw-Sentinel-2 GAN with perceptual + spectral-norm + BCE, stronger GAN. + +Same three additions as ``naip_gan_s2_percep_sn_bce.py`` (VGG perceptual loss, +spectrally-normalized discriminator, BCE adversarial loss) but pushes the +adversarial objective harder now that spectral norm stabilizes the +discriminator: the adversarial weight is raised (``lambda_adv=0.5``) while L1 +and perceptual stay at 1.0, and the adversarial warmup is shortened so the GAN +term engages sooner. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_percep_sn_bce_moregan.py dry_run naip_gan_s2_percep_sn_bce_moregan local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2_percep_sn_bce as percep_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +LAMBDA_ADV = 0.5 +GAN_WARMUP_STEPS = 2000 + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the perceptual + SN + BCE config with a stronger adversarial term.""" + config = percep_base.build_train_module_config(common) + config.lambda_adv = LAMBDA_ADV + config.gan_warmup_steps = GAN_WARMUP_STEPS + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_no_l1.py b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_no_l1.py new file mode 100644 index 000000000..884a015e8 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_percep_sn_bce_no_l1.py @@ -0,0 +1,48 @@ +"""v1.2 NAIP raw-Sentinel-2 GAN with perceptual as the sole content anchor. + +Same three additions as ``naip_gan_s2_percep_sn_bce.py`` (VGG perceptual loss, +spectrally-normalized discriminator, BCE adversarial loss) but drops the L1 +pixel loss (``lambda_l1=0``). This is the corrected "no-L1" recipe: unlike a +bare adversarial-only run, the perceptual loss still supplies a strong +content/alignment gradient (exactly what ESRGAN keeps when its L1 term is +removed), so the generator has a reconstruction anchor. A short adversarial +warmup lets the perceptual term shape the output first. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_percep_sn_bce_no_l1.py dry_run naip_gan_s2_percep_sn_bce_no_l1 local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2_percep_sn_bce as percep_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +# Perceptual replaces L1 as the content anchor; keep the adversarial weight at +# the ESRGAN reference. A short warmup lets the perceptual term act first. +GAN_WARMUP_STEPS = 2000 + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the perceptual + SN + BCE config with L1 disabled.""" + config = percep_base.build_train_module_config(common) + config.lambda_l1 = 0.0 + config.gan_warmup_steps = GAN_WARMUP_STEPS + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_projdis.py b/scripts/official/v1_2/naip_gan_s2_projdis.py new file mode 100644 index 000000000..0a236357e --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_projdis.py @@ -0,0 +1,38 @@ +"""v1.2 NAIP raw-Sentinel-2 GAN with a projection discriminator. + +Same as ``naip_gan_s2.py`` (raw Sentinel-2 condition) but enables the +projection-discriminator term (``use_projection=True``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_projdis.py dry_run naip_gan_s2_projdis local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2 as naip_gan_s2_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the raw-Sentinel-2 train module config with the projection term.""" + config = naip_gan_s2_base.build_train_module_config(common) + config.discriminator_config.use_projection = True + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_projdis_no_l1.py b/scripts/official/v1_2/naip_gan_s2_projdis_no_l1.py new file mode 100644 index 000000000..0018de938 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_projdis_no_l1.py @@ -0,0 +1,43 @@ +"""v1.2 NAIP raw-Sentinel-2 projection GAN without the L1 loss. + +Same as ``naip_gan_s2_projdis.py`` (raw Sentinel-2 condition, projection +discriminator) but drops the L1 reconstruction loss (``lambda_l1=0``), so the +generator is trained by the adversarial loss alone (``lambda_adv=0.1``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_projdis_no_l1.py dry_run naip_gan_s2_projdis_no_l1 local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2 as naip_gan_s2_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the raw-Sentinel-2 train module config: projection, no L1 loss.""" + config = naip_gan_s2_base.build_train_module_config(common) + config.discriminator_config.use_projection = True + config.lambda_l1 = 0.0 + # No L1 term, so there is nothing to train the generator during a warmup; + # start the adversarial loss immediately. + config.gan_warmup_steps = 0 + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_s2_projdis_percep_sn_bce.py b/scripts/official/v1_2/naip_gan_s2_projdis_percep_sn_bce.py new file mode 100644 index 000000000..b83a5b1e4 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_s2_projdis_percep_sn_bce.py @@ -0,0 +1,41 @@ +"""v1.2 NAIP raw-Sentinel-2 projection GAN with perceptual + spectral-norm + BCE. + +Same three additions as ``naip_gan_s2_percep_sn_bce.py`` (VGG perceptual loss, +spectrally-normalized discriminator, BCE adversarial loss) but with the +projection-discriminator term enabled (``use_projection=True``). Tests whether +the projection discriminator helps once it is stabilized by spectral norm. Loss +balance matches the ESRGAN reference (L1=1.0, perceptual=1.0, adv=0.1). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_s2_projdis_percep_sn_bce.py dry_run naip_gan_s2_projdis_percep_sn_bce local +""" + +import logging + +import naip_gan as naip_gan_base +import naip_gan_s2_percep_sn_bce as percep_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the perceptual + SN + BCE config with the projection term on.""" + config = percep_base.build_train_module_config(common) + config.discriminator_config.use_projection = True + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_unmasked_embed.py b/scripts/official/v1_2/naip_gan_unmasked_embed.py new file mode 100644 index 000000000..5fa568491 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_unmasked_embed.py @@ -0,0 +1,33 @@ +"""v1.2 NAIP conditional GAN conditioned on the unmasked-input embedding. + +This is the base ``naip_gan.py`` configuration made explicit: the discriminator +conditions on the online encoder's pooled embedding from a full-depth forward on +the unmasked input (``online_unmasked_pooled``) with the ``[256, 128, 128]`` +generator and the condition-aware discriminator (two head residual blocks, +post-unpatchify condition convs, no projection term). It exists as a named +experiment; it reuses the base builders unchanged. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_unmasked_embed.py dry_run naip_gan_unmasked_embed local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.internal.experiment import main + +logger = logging.getLogger(__name__) + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=naip_gan_base.build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_unmasked_embed_moregan.py b/scripts/official/v1_2/naip_gan_unmasked_embed_moregan.py new file mode 100644 index 000000000..a0d23149a --- /dev/null +++ b/scripts/official/v1_2/naip_gan_unmasked_embed_moregan.py @@ -0,0 +1,44 @@ +"""v1.2 NAIP unmasked-embedding conditional GAN with a stronger adversarial term. + +Same as ``naip_gan_unmasked_embed.py`` (discriminator conditions on the online +encoder's pooled embedding from a full-depth forward on the unmasked input) but +rebalances the losses toward the GAN objective: the L1 reconstruction weight is +lowered (``lambda_l1=2.0``) and the adversarial weight is raised +(``lambda_adv=0.5``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_unmasked_embed_moregan.py dry_run naip_gan_unmasked_embed_moregan local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +LAMBDA_L1 = 2.0 +LAMBDA_ADV = 0.5 + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the unmasked-embedding train module config with a stronger GAN term.""" + config = naip_gan_base.build_train_module_config(common) + config.lambda_l1 = LAMBDA_L1 + config.lambda_adv = LAMBDA_ADV + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_unmasked_embed_projdis.py b/scripts/official/v1_2/naip_gan_unmasked_embed_projdis.py new file mode 100644 index 000000000..c84f70896 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_unmasked_embed_projdis.py @@ -0,0 +1,38 @@ +"""v1.2 NAIP unmasked-embedding GAN with a projection discriminator. + +Same as ``naip_gan_unmasked_embed.py`` but enables the projection-discriminator +term (``use_projection=True``), so the discriminator's logits reward agreement +between the NAIP image features and the projected condition. + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_unmasked_embed_projdis.py dry_run naip_gan_unmasked_embed_projdis local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the train module config with the projection-discriminator term.""" + config = naip_gan_base.build_train_module_config(common) + config.discriminator_config.use_projection = True + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_gan_unmasked_embed_projdis_no_l1.py b/scripts/official/v1_2/naip_gan_unmasked_embed_projdis_no_l1.py new file mode 100644 index 000000000..43aa02603 --- /dev/null +++ b/scripts/official/v1_2/naip_gan_unmasked_embed_projdis_no_l1.py @@ -0,0 +1,43 @@ +"""v1.2 NAIP unmasked-embedding projection GAN without the L1 loss. + +Same as ``naip_gan_unmasked_embed_projdis.py`` (unmasked-input embedding +condition, projection discriminator) but drops the L1 reconstruction loss +(``lambda_l1=0``), so the generator is trained by the adversarial loss alone +(``lambda_adv=0.1``). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_gan_unmasked_embed_projdis_no_l1.py dry_run naip_gan_unmasked_embed_projdis_no_l1 local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the train module config: projection discriminator, no L1 loss.""" + config = naip_gan_base.build_train_module_config(common) + config.discriminator_config.use_projection = True + config.lambda_l1 = 0.0 + # No L1 term, so there is nothing to train the generator during a warmup; + # start the adversarial loss immediately. + config.gan_warmup_steps = 0 + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=naip_gan_base.build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/scripts/official/v1_2/naip_l1_only.py b/scripts/official/v1_2/naip_l1_only.py new file mode 100644 index 000000000..5bdf7b86a --- /dev/null +++ b/scripts/official/v1_2/naip_l1_only.py @@ -0,0 +1,64 @@ +"""v1.2 NAIP reconstruction with L1 loss only (no GAN). + +The NAIP GAN experiments (``naip_gan.py`` and its variants) were not converging, +so this config drops the adversarial branch entirely: the NAIP generator on top +of the encoder's pooled spatial embedding is trained by the L1 reconstruction +loss alone (``lambda_adv=0``, no discriminator built, no perceptual loss). The +latent-MIM patch-discrimination SSL objective for the encoder is unchanged. + +The generator is also slimmed from the base ``[256, 192, 192]`` per-stage widths +(~11M params) to ``[192, 160, 128]`` (~7M params). + +Validate before launching:: + + python3 scripts/official/v1_2/naip_l1_only.py dry_run naip_l1_only local +""" + +import logging + +import naip_gan as naip_gan_base + +from olmoearth_pretrain.internal.experiment import CommonComponents, main +from olmoearth_pretrain.nn.naip_gan import NaipGanModelConfig +from olmoearth_pretrain.train.train_module.naip_gan import NaipGanTrainModuleConfig + +logger = logging.getLogger(__name__) + +# Slimmer generator (base resolution + one per upsampling stage). Lands at +# ~7M params vs. the base ``[256, 192, 192]`` (~11M), within the 5-10M target. +GENERATOR_HIDDEN_SIZES = [192, 160, 128] + + +def build_model_config(common: CommonComponents) -> NaipGanModelConfig: + """Build the NAIP model config with a slimmed generator (no discriminator).""" + config = naip_gan_base.build_model_config(common) + config.generator_config.hidden_sizes = GENERATOR_HIDDEN_SIZES + return config + + +def build_train_module_config(common: CommonComponents) -> NaipGanTrainModuleConfig: + """Build the train module config: L1 reconstruction only, no GAN.""" + config = naip_gan_base.build_train_module_config(common) + # Drop the adversarial branch: no discriminator (or its optimizer) is built, + # and the generator is trained by L1 alone from step 0. + config.lambda_adv = 0.0 + config.lambda_perceptual = 0.0 + # L1 is the only NAIP loss now, so it no longer has to out-weigh an + # adversarial term; lower it from the base 10.0 to 2.0. + config.lambda_l1 = 2.0 + config.discriminator_config = None + config.disc_optim_config = None + config.gan_warmup_steps = 0 + return config + + +if __name__ == "__main__": + main( + common_components_builder=naip_gan_base.build_common_components, + model_config_builder=build_model_config, + train_module_config_builder=build_train_module_config, + dataset_config_builder=naip_gan_base.build_dataset_config, + dataloader_config_builder=naip_gan_base.build_dataloader_config, + trainer_config_builder=naip_gan_base.build_trainer_config, + visualize_config_builder=naip_gan_base.build_visualize_config, + ) diff --git a/tests/integration/train/train_module/test_naip_gan.py b/tests/integration/train/train_module/test_naip_gan.py new file mode 100644 index 000000000..5b8657b7b --- /dev/null +++ b/tests/integration/train/train_module/test_naip_gan.py @@ -0,0 +1,443 @@ +"""Integration tests for the NAIP GAN train module. + +These validate the gradient isolation that defines the objective: +- the discriminator loss updates only the discriminator (not the encoder), and +- the generator loss updates the generator/encoder (not the discriminator). +""" + +import logging +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from olmo_core.optim.adamw import AdamWConfig + +from olmoearth_pretrain.data.collate import collate_single_masked_batched +from olmoearth_pretrain.data.constants import Modality +from olmoearth_pretrain.data.dataset import OlmoEarthSample +from olmoearth_pretrain.data.transform import TransformConfig +from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample +from olmoearth_pretrain.nn.flexi_vit import EncoderConfig, PredictorConfig +from olmoearth_pretrain.nn.naip_gan import ( + NaipDiscriminatorConfig, + NaipGanModel, + NaipGanModelConfig, + NaipGeneratorConfig, +) +from olmoearth_pretrain.train.loss import LossConfig +from olmoearth_pretrain.train.masking import MaskingConfig +from olmoearth_pretrain.train.train_module.naip_gan import ( + NaipGanTrainModule, + NaipGanTrainModuleConfig, +) +from olmoearth_pretrain.train.utils import split_masked_batch + +torch.set_default_device("cpu") +logger = logging.getLogger(__name__) + +SUPPORTED_MODALITIES = [ + Modality.SENTINEL2_L2A.name, + Modality.SENTINEL1.name, + Modality.WORLDCOVER.name, + Modality.LATLON.name, + Modality.NAIP_10.name, +] +EMBEDDING_SIZE = 16 + + +class MockTrainer: + """Minimal trainer stub for driving a single train batch.""" + + def __init__(self) -> None: + """Initialize the mock trainer.""" + self._metrics: dict[str, float] = {} + self.global_step = 0 + self.max_steps = 100 + + def record_metric( + self, + name: str, + value: float, + reduce_type: str, + namespace: str | None = None, + ) -> None: + """Record a metric.""" + self._metrics[name] = value + + def _iter_callbacks(self): # type: ignore[no-untyped-def] + """No callbacks in the mock trainer.""" + return iter([]) + + +@pytest.fixture +def naip_samples() -> list[tuple[int, OlmoEarthSample]]: + """Two samples with a present (decode-only) NAIP modality.""" + s2 = np.random.randn(8, 8, 12, 13).astype(np.float32) + s1 = np.random.randn(8, 8, 12, 2).astype(np.float32) + wc = np.random.randn(8, 8, 1, 10).astype(np.float32) + # naip_10 spatial extent is height * image_tile_size_factor (4) -> 32x32, 4 bands. + naip = np.random.randn(32, 32, 1, 4).astype(np.float32) + latlon = np.random.randn(2).astype(np.float32) + timestamps = np.array( + [ + [15, 7, 2023], + [15, 8, 2023], + [15, 9, 2023], + [15, 10, 2023], + [15, 11, 2023], + [15, 11, 2023], + [15, 1, 2024], + [15, 2, 2024], + [15, 3, 2024], + [15, 4, 2024], + [15, 5, 2024], + [15, 6, 2024], + ], + dtype=np.int32, + ) + + def make() -> OlmoEarthSample: + return OlmoEarthSample( + sentinel2_l2a=s2, + sentinel1=s1, + worldcover=wc, + naip_10=naip, + latlon=latlon, + timestamps=timestamps, + ) + + return [(1, make()), (1, make())] + + +@pytest.fixture +def naip_gan_model() -> NaipGanModel: + """A tiny NAIP GAN model on CPU.""" + encoder_config = EncoderConfig( + supported_modality_names=SUPPORTED_MODALITIES, + embedding_size=EMBEDDING_SIZE, + min_patch_size=1, + max_patch_size=1, + num_heads=2, + mlp_ratio=1.0, + depth=2, + drop_path=0.0, + max_sequence_length=12, + ) + decoder_config = PredictorConfig( + supported_modality_names=SUPPORTED_MODALITIES, + encoder_embedding_size=EMBEDDING_SIZE, + decoder_embedding_size=EMBEDDING_SIZE, + depth=2, + mlp_ratio=1.0, + num_heads=2, + max_sequence_length=12, + drop_path=0.0, + output_embedding_size=None, + ) + generator_config = NaipGeneratorConfig( + embedding_size=EMBEDDING_SIZE, + patch_size=1, + hidden_sizes=[32, 32, 32], + out_channels=4, + upsample_factor=4, + num_res_blocks=1, + ) + model = NaipGanModelConfig( + encoder_config=encoder_config, + decoder_config=decoder_config, + generator_config=generator_config, + ).build() + model.to(device="cpu") + return model + + +def _build_train_module( + model: NaipGanModel, + image_log_interval: int = 0, + cond_source: str = "online_pooled", +) -> NaipGanTrainModule: + if cond_source == "raw_sentinel2": + # Condition on the full Sentinel-2 temporal stack (all 13 sample bands). + discriminator_config = NaipDiscriminatorConfig( + embedding_size=EMBEDDING_SIZE, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_mode="image", + cond_in_channels=13, + ) + else: + discriminator_config = NaipDiscriminatorConfig( + embedding_size=EMBEDDING_SIZE, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + ) + config = NaipGanTrainModuleConfig( + optim_config=AdamWConfig(lr=1e-4, weight_decay=0.0), + rank_microbatch_size=2, + loss_config=LossConfig(loss_config={"type": "patch_discrimination"}), + masking_config=MaskingConfig( + strategy_config={ + "type": "random_time_with_decode", + "encode_ratio": 0.5, + "decode_ratio": 0.5, + "random_ratio": 0.5, + "only_decode_modalities": [Modality.NAIP_10.name], + } + ), + discriminator_config=discriminator_config, + disc_optim_config=AdamWConfig(lr=2e-4, betas=(0.5, 0.999), weight_decay=0.0), + lambda_adv=1.0, + lambda_l1=1.0, + gan_warmup_steps=0, + image_log_interval=image_log_interval, + discriminator_cond_source=cond_source, + token_exit_cfg={modality: 0 for modality in Modality.names()}, + ema_decay=(1.0, 1.0), + max_grad_norm=1.0, + transform_config=TransformConfig(transform_type="no_transform"), + ) + tm = config.build(model, device="cpu") + tm.on_attach = MagicMock(return_value=None) # type: ignore + tm._attach_trainer(MockTrainer()) + return tm + + +def _collate( + naip_samples: list[tuple[int, OlmoEarthSample]], +) -> tuple[int, MaskedOlmoEarthSample]: + masking_strategy = MaskingConfig( + strategy_config={ + "type": "random_time_with_decode", + "encode_ratio": 0.5, + "decode_ratio": 0.5, + "random_ratio": 0.5, + "only_decode_modalities": [Modality.NAIP_10.name], + } + ).build() + return collate_single_masked_batched( + naip_samples, transform=None, masking_strategy=masking_strategy + ) + + +def _has_any_grad(module: torch.nn.Module) -> bool: + return any( + p.grad is not None and torch.any(p.grad != 0) for p in module.parameters() + ) + + +def _all_grads_none(module: torch.nn.Module) -> bool: + return all(p.grad is None for p in module.parameters()) + + +def test_train_batch_runs_and_records_metrics( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """A full train batch runs end-to-end and records the GAN metrics.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_train_module(naip_gan_model) + tm.train_batch(batch) + assert "train/G_l1" in tm.trainer._metrics + assert "train/D_loss" in tm.trainer._metrics + # Generator (in the model) received gradients from the GAN branch. + assert _has_any_grad(tm.model.generator) + + +def test_gradient_isolation( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """D loss updates only D; G loss updates only the generator/encoder.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_train_module(naip_gan_model) + patch_size, batch_data = batch + microbatch = split_masked_batch(batch_data, tm.rank_microbatch_size)[0] + microbatch = microbatch.to_device(torch.device("cpu")) + + _, _, _, _, pooled, fake_naip = tm.model_forward( + microbatch, patch_size, tm.token_exit_cfg + ) + cond, cond_valid, cond_time_mask = tm._extract_discriminator_cond( + microbatch, pooled, patch_size + ) + + # --- Discriminator loss: only the discriminator gets gradients. --- + tm.zero_grads() + d_loss = tm._maybe_discriminator_loss( + microbatch, + cond, + cond_valid, + fake_naip, + num_microbatches=1, + adversarial_active=True, + patch_size=patch_size, + cond_time_mask=cond_time_mask, + ) + assert d_loss is not None + d_loss.backward() + assert _has_any_grad(tm.discriminator) + assert _all_grads_none(tm.model.encoder) + assert _all_grads_none(tm.model.generator) + + # --- Generator loss: only the generator/encoder get gradients. --- + tm.zero_grads() + gen_loss, _, _, _ = tm._generator_loss( + microbatch, + cond, + cond_valid, + fake_naip, + adversarial_active=True, + patch_size=patch_size, + cond_time_mask=cond_time_mask, + ) + gen_loss.backward() + assert _all_grads_none(tm.discriminator) + assert _has_any_grad(tm.model.generator) + assert _has_any_grad(tm.model.encoder) + + +@pytest.mark.parametrize( + "cond_source", ["online_pooled", "online_unmasked_pooled", "raw_sentinel2"] +) +def test_train_batch_runs_for_cond_sources( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, + cond_source: str, +) -> None: + """Each discriminator conditioning source runs a full batch end-to-end.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_train_module(naip_gan_model, cond_source=cond_source) + tm.train_batch(batch) + assert "train/G_l1" in tm.trainer._metrics + assert "train/D_loss" in tm.trainer._metrics + # The discriminator was trained (D loss active from step 0 here). + assert _has_any_grad(tm.discriminator) + # Generator (in the model) received gradients from the GAN branch. + assert _has_any_grad(tm.model.generator) + + +def _build_l1_only_train_module(model: NaipGanModel) -> NaipGanTrainModule: + """Build a train module with no discriminator (pure L1 reconstruction).""" + config = NaipGanTrainModuleConfig( + optim_config=AdamWConfig(lr=1e-4, weight_decay=0.0), + rank_microbatch_size=2, + loss_config=LossConfig(loss_config={"type": "patch_discrimination"}), + masking_config=MaskingConfig( + strategy_config={ + "type": "random_time_with_decode", + "encode_ratio": 0.5, + "decode_ratio": 0.5, + "random_ratio": 0.5, + "only_decode_modalities": [Modality.NAIP_10.name], + } + ), + discriminator_config=None, + disc_optim_config=None, + lambda_adv=0.0, + lambda_l1=1.0, + gan_warmup_steps=0, + image_log_interval=0, + token_exit_cfg={modality: 0 for modality in Modality.names()}, + ema_decay=(1.0, 1.0), + max_grad_norm=1.0, + transform_config=TransformConfig(transform_type="no_transform"), + ) + tm = config.build(model, device="cpu") + tm.on_attach = MagicMock(return_value=None) # type: ignore + tm._attach_trainer(MockTrainer()) + return tm + + +def test_l1_only_requires_no_discriminator( + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """With lambda_adv=0 the module builds without a discriminator/optimizer.""" + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_l1_only_train_module(naip_gan_model) + assert tm.discriminator is None + assert tm.disc_optimizer is None + + +def test_l1_only_train_batch_trains_generator( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """A pure-L1 batch runs end-to-end, trains the generator, records no D loss.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_l1_only_train_module(naip_gan_model) + tm.train_batch(batch) + # L1 is recorded; the adversarial/discriminator metrics are not. + assert "train/G_l1" in tm.trainer._metrics + assert "train/G_adv" not in tm.trainer._metrics + assert "train/D_loss" not in tm.trainer._metrics + # The generator (and encoder) are trained by the L1 loss alone. + assert _has_any_grad(tm.model.generator) + + +class FakeWandb: + """Minimal wandb stand-in capturing Image/log calls.""" + + def __init__(self) -> None: + """Initialize the fake wandb module.""" + self.logged: list[tuple[dict, int | None]] = [] + + def Image(self, arr: np.ndarray, caption: str | None = None) -> np.ndarray: # noqa: N802 + """Return the array unchanged (records nothing extra).""" + return arr + + def log(self, data: dict, step: int | None = None) -> None: + """Capture a log call.""" + self.logged.append((data, step)) + + +def test_naip_images_logged_at_interval( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """Paired real/generated NAIP images are logged to W&B at the interval.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_train_module(naip_gan_model, image_log_interval=1) + fake_wandb = FakeWandb() + tm._get_wandb = lambda: fake_wandb # type: ignore[method-assign] + tm.train_batch(batch) + + assert len(fake_wandb.logged) == 1 + data, step = fake_wandb.logged[0] + assert step == tm.trainer.global_step + images = data["gan/naip_real_vs_fake"] + # Two valid NAIP samples in the batch -> two paired images. + assert len(images) == 2 + for pair in images: + # Side-by-side [H, 2W, 3] RGB in [0, 1]. + assert pair.ndim == 3 + assert pair.shape[2] == 3 + assert float(pair.min()) >= 0.0 and float(pair.max()) <= 1.0 + + +def test_naip_images_not_logged_when_disabled( + naip_samples: list[tuple[int, OlmoEarthSample]], + naip_gan_model: NaipGanModel, + set_random_seeds: None, +) -> None: + """No images are logged when image_log_interval is 0.""" + batch = _collate(naip_samples) + with patch("olmoearth_pretrain.train.train_module.train_module.build_world_mesh"): + tm = _build_train_module(naip_gan_model, image_log_interval=0) + fake_wandb = FakeWandb() + tm._get_wandb = lambda: fake_wandb # type: ignore[method-assign] + tm.train_batch(batch) + assert len(fake_wandb.logged) == 0 diff --git a/tests/unit/nn/test_naip_gan.py b/tests/unit/nn/test_naip_gan.py new file mode 100644 index 000000000..5d9e6cd54 --- /dev/null +++ b/tests/unit/nn/test_naip_gan.py @@ -0,0 +1,354 @@ +"""Unit tests for the NAIP GAN generator and discriminator.""" + +import pytest +import torch + +from olmoearth_pretrain.nn.naip_gan import ( + NaipDiscriminator, + NaipGenerator, + discriminator_adversarial_loss, + generator_adversarial_loss, +) + +torch.set_default_device("cpu") + + +def test_generator_output_shape() -> None: + """Generator upsamples the token grid by patch_size * upsample_factor.""" + batch, height, width, dim = 2, 4, 5, 16 + patch_size, upsample_factor = 4, 4 + gen = NaipGenerator( + embedding_size=dim, + patch_size=patch_size, + hidden_sizes=[32, 32, 32], + out_channels=4, + upsample_factor=upsample_factor, + num_res_blocks=1, + ) + pooled = torch.randn(batch, height, width, dim) + out = gen(pooled, patch_size=patch_size) + factor = patch_size * upsample_factor + assert out.shape == (batch, 4, height * factor, width * factor) + + +def test_generator_upsample_factor_must_be_power_of_two() -> None: + """A non-power-of-2 upsample factor is rejected.""" + with pytest.raises(ValueError): + NaipGenerator( + embedding_size=8, patch_size=4, hidden_sizes=[8, 8, 8], upsample_factor=3 + ) + + +@pytest.mark.parametrize("in_height,in_width", [(4, 5), (2, 2), (7, 3)]) +def test_generator_resamples_up_for_large_patch_size( + in_height: int, in_width: int +) -> None: + """A larger encoder patch size upsamples the token grid before unpatchify.""" + dim = 16 + unpatchify, upsample = 4, 4 + gen = NaipGenerator( + embedding_size=dim, + patch_size=unpatchify, # canonical unpatchify factor U + hidden_sizes=[32, 32, 32], + out_channels=4, + upsample_factor=upsample, + num_res_blocks=1, + ) + pooled = torch.randn(2, in_height, in_width, dim) + # patch_size == U -> no resample. + out = gen(pooled, patch_size=unpatchify) + assert out.shape == ( + 2, + 4, + in_height * unpatchify * upsample, + in_width * unpatchify * upsample, + ) + # patch_size 8 (> U=4) -> resample the token grid by 8/4 = 2. + out2 = gen(pooled, patch_size=8) + assert out2.shape == ( + 2, + 4, + in_height * 2 * unpatchify * upsample, + in_width * 2 * unpatchify * upsample, + ) + + +def test_generator_resamples_down_for_small_patch_size() -> None: + """A smaller encoder patch size downsamples the token grid before unpatchify.""" + dim = 8 + gen = NaipGenerator( + embedding_size=dim, + patch_size=4, + hidden_sizes=[16, 16], + upsample_factor=2, + num_res_blocks=1, + ) + pooled = torch.randn(1, 8, 8, dim) + # patch_size 2 (< U=4) -> resample by 2/4 = 0.5 -> 4x4 tokens. + out = gen(pooled, patch_size=2) + factor = 4 * 2 # U * upsample_factor + assert out.shape == (1, 4, 4 * factor, 4 * factor) + + +def test_generator_canonical_patch_size_uses_token_grid() -> None: + """When patch_size equals the unpatchify factor the token grid is used as-is.""" + dim = 8 + gen = NaipGenerator( + embedding_size=dim, + patch_size=2, + hidden_sizes=[16, 16], + upsample_factor=2, + num_res_blocks=1, + ) + out = gen(torch.randn(1, 4, 4, dim), patch_size=2) + factor = 2 * 2 + assert out.shape == (1, 4, 4 * factor, 4 * factor) + + +def test_discriminator_output_shape() -> None: + """Discriminator returns per-patch logits at the conditioning grid size.""" + batch, height, width, dim = 2, 4, 5, 16 + disc = NaipDiscriminator( + embedding_size=dim, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + ) + image = torch.randn(batch, 4, height * 4, width * 4) + cond = torch.randn(batch, height, width, dim) + logits = disc(image, cond, patch_size=1) + assert logits.shape == (batch, 1, height, width) + + +def test_generator_casts_input_to_param_dtype() -> None: + """Generator handles inputs whose dtype differs from its params (mixed precision).""" + gen = NaipGenerator( + embedding_size=8, + patch_size=2, + hidden_sizes=[16, 16], + upsample_factor=2, + num_res_blocks=1, + ).double() + pooled = torch.randn(1, 4, 4, 8) # float32 input, float64 params + out = gen(pooled, patch_size=2) + assert out.dtype == torch.float64 + + +def test_discriminator_casts_inputs_to_param_dtype() -> None: + """Discriminator handles inputs whose dtype differs from its params.""" + disc = NaipDiscriminator( + embedding_size=8, + in_channels=4, + image_strided_conv_channels=[8, 16], + feature_channels=16, + ).double() + image = torch.randn(1, 4, 16, 16) # float32 inputs, float64 params + cond = torch.randn(1, 4, 4, 8) + logits = disc(image, cond, patch_size=1) + assert logits.dtype == torch.float64 + + +def test_discriminator_image_cond_output_shape() -> None: + """Image-conditioned discriminator handles a Sentinel-2 temporal stack.""" + batch, t, grid_h, grid_w = 2, 5, 4, 5 + disc = NaipDiscriminator( + embedding_size=16, # unused in image mode + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_mode="image", + cond_in_channels=6, + ) + image = torch.randn(batch, 4, grid_h * 4, grid_w * 4) + # Temporal stack [B, T, C, H, W] at a Sentinel-2-like resolution; fusion + # happens at that image resolution. + cond_h, cond_w = grid_h * 2, grid_w * 2 + cond = torch.randn(batch, t, 6, cond_h, cond_w) + time_mask = torch.ones(batch, t, dtype=torch.bool) + logits = disc(image, cond, patch_size=1, cond_time_mask=time_mask) + assert logits.shape == (batch, 1, cond_h, cond_w) + + +def test_discriminator_unknown_cond_mode_raises() -> None: + """An unsupported cond_mode is rejected.""" + with pytest.raises(ValueError): + NaipDiscriminator(embedding_size=8, cond_mode="nope") + + +def test_discriminator_cond_pre_post_pool_channels() -> None: + """The S2 cond stem uses non-strided pre-pool and post-pool conv stacks.""" + batch, t, grid_h, grid_w = 2, 4, 4, 5 + disc = NaipDiscriminator( + embedding_size=8, # unused in image mode + in_channels=4, + image_strided_conv_channels=[16, 32, 64], + feature_channels=64, + cond_mode="image", + cond_in_channels=6, + cond_image_pre_pool_channels=[16, 24], + cond_image_post_pool_channels=[24], + ) + image = torch.randn(batch, 4, grid_h * 8, grid_w * 8) + cond_h, cond_w = grid_h * 2, grid_w * 2 + cond = torch.randn(batch, t, 6, cond_h, cond_w) + logits = disc(image, cond, patch_size=1) + assert logits.shape == (batch, 1, cond_h, cond_w) + + +def test_discriminator_image_cond_no_time_mask() -> None: + """Without a time mask the condition averages over all timesteps.""" + batch, t, grid_h, grid_w = 2, 3, 4, 4 + disc = NaipDiscriminator( + embedding_size=8, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_mode="image", + cond_in_channels=6, + ) + image = torch.randn(batch, 4, grid_h * 4, grid_w * 4) + cond_h, cond_w = grid_h * 3, grid_w * 3 + cond = torch.randn(batch, t, 6, cond_h, cond_w) + logits = disc(image, cond, patch_size=1) + assert logits.shape == (batch, 1, cond_h, cond_w) + + +def test_discriminator_cond_embedding_channels() -> None: + """cond_embedding_channels adds post-unpatchify convs to the condition path.""" + disc = NaipDiscriminator( + embedding_size=16, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_embedding_channels=[64, 32], + ) + # cond_proj is a single linear (the unpatchify); the convs live in + # cond_post_unpatchify. + assert isinstance(disc.cond_proj, torch.nn.Linear) + num_convs = sum( + 1 for m in disc.cond_post_unpatchify if isinstance(m, torch.nn.Conv2d) + ) + assert num_convs == 3 # two 3x3 convs (64, 32) plus the final 1x1 projection + image = torch.randn(2, 4, 16, 20) + cond = torch.randn(2, 4, 5, 16) + logits = disc(image, cond, patch_size=1) + assert logits.shape == (2, 1, 4, 5) + + +def test_discriminator_no_cond_embedding_channels_is_identity() -> None: + """Without cond_embedding_channels there are no post-unpatchify convs.""" + disc = NaipDiscriminator( + embedding_size=16, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + ) + num_convs = sum( + 1 for m in disc.cond_post_unpatchify if isinstance(m, torch.nn.Conv2d) + ) + assert num_convs == 0 + + +def test_discriminator_convs_per_resolution() -> None: + """Refinement convs add layers without changing the output shape.""" + batch, height, width, dim = 2, 4, 5, 16 + disc = NaipDiscriminator( + embedding_size=dim, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + num_convs_per_resolution=2, + ) + # stem + 1 strided conv, each followed by 2 refinement convs (6), plus the + # final projection conv to feature_channels (1) -> 7 convs total. + num_convs = sum(1 for m in disc.from_image if isinstance(m, torch.nn.Conv2d)) + assert num_convs == 7 + image = torch.randn(batch, 4, height * 4, width * 4) + cond = torch.randn(batch, height, width, dim) + logits = disc(image, cond, patch_size=1) + assert logits.shape == (batch, 1, height, width) + + +def test_discriminator_cond_unpatchify_factor_embedding() -> None: + """Learned unpatchify fuses at the base token_grid * patch_size grid.""" + batch, height, width, dim = 2, 4, 5, 16 + f = 4 + disc = NaipDiscriminator( + embedding_size=dim, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + use_projection=True, + cond_unpatchify_factor=f, + ) + image = torch.randn(batch, 4, height * 4 * f, width * 4 * f) + cond = torch.randn(batch, height, width, dim) + # patch_size == f: tokens already canonical, fusion grid = token_grid * f. + logits = disc(image, cond, patch_size=f) + assert logits.shape == (batch, 1, height * f, width * f) + + +def test_discriminator_cond_unpatchify_factor_image() -> None: + """Image cond mode fuses at the Sentinel-2 image resolution.""" + batch, t, grid_h, grid_w = 2, 4, 3, 4 + disc = NaipDiscriminator( + embedding_size=8, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_mode="image", + cond_in_channels=6, + ) + image = torch.randn(batch, 4, grid_h * 8, grid_w * 8) + cond_h, cond_w = grid_h * 4, grid_w * 4 + cond = torch.randn(batch, t, 6, cond_h, cond_w) + time_mask = torch.ones(batch, t, dtype=torch.bool) + logits = disc(image, cond, patch_size=1, cond_time_mask=time_mask) + assert logits.shape == (batch, 1, cond_h, cond_w) + + +def test_discriminator_invalid_cond_unpatchify_factor() -> None: + """cond_unpatchify_factor must be >= 1.""" + with pytest.raises(ValueError): + NaipDiscriminator(embedding_size=8, cond_unpatchify_factor=0) + + +def test_discriminator_resamples_cond_for_patch_size() -> None: + """A patch size below the unpatchify factor resamples the condition down.""" + batch, height, width, dim = 2, 8, 8, 16 + f = 4 + disc = NaipDiscriminator( + embedding_size=dim, + in_channels=4, + image_strided_conv_channels=[16, 32], + feature_channels=32, + cond_unpatchify_factor=f, + ) + cond = torch.randn(batch, height, width, dim) + image = torch.randn(batch, 4, 64, 64) + # patch_size 2 (< f=4): tokens -> round(8 * 2 / 4) = 4, fusion grid = 4 * f. + logits = disc(image, cond, patch_size=2) + assert logits.shape == (batch, 1, 4 * f, 4 * f) + + +@pytest.mark.parametrize("loss_type", ["hinge", "bce"]) +def test_adversarial_losses_are_scalars(loss_type: str) -> None: + """Both adversarial losses reduce to scalars for each variant.""" + real_logits = torch.randn(2, 1, 4, 4, requires_grad=True) + fake_logits = torch.randn(2, 1, 4, 4, requires_grad=True) + d_loss = discriminator_adversarial_loss(real_logits, fake_logits, loss_type) + g_loss = generator_adversarial_loss(fake_logits, loss_type) + assert d_loss.ndim == 0 + assert g_loss.ndim == 0 + # Losses should be differentiable. + d_loss.backward(retain_graph=True) + g_loss.backward() + + +def test_unknown_loss_type_raises() -> None: + """An unsupported loss type is rejected.""" + logits = torch.randn(2, 1, 4, 4) + with pytest.raises(ValueError): + generator_adversarial_loss(logits, "nope") + with pytest.raises(ValueError): + discriminator_adversarial_loss(logits, logits, "nope")