# Copyright 2025 The Galileo Authors and The HuggingFace Inc. team. """Self-contained Galileo model, config, and encoder implementation.""" from __future__ import annotations from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig from transformers.modeling_outputs import BaseModelOutputWithPooling from transformers.modeling_utils import PreTrainedModel from transformers.processing_utils import Unpack from transformers.utils import TransformersKwargs, logging import collections.abc import itertools import json import math from collections import OrderedDict from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple, Union from typing import OrderedDict as OrderedDictType import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor, vmap from torch.jit import Final def _band_group_indices( bands: Sequence[str], groups: Dict[str, List[str]] ) -> OrderedDictType[str, List[int]]: return OrderedDict((name, [bands.index(b) for b in group_bands]) for name, group_bands in groups.items()) def get_2d_sincos_pos_embed_with_resolution( embed_dim, grid_size, res, cls_token=False, device="cpu" ): """ grid_size: int of the grid height and width res: array of size n, representing the resolution of a pixel (say, in meters), return: pos_embed: [n,grid_size*grid_size, embed_dim] or [n,1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) """ res = res.to(device) grid_h = torch.arange(grid_size, device=device) grid_w = torch.arange(grid_size, device=device) grid = torch.meshgrid( grid_w, grid_h, indexing="xy" ) # here h goes first,direction reversed for numpy grid = torch.stack(grid, dim=0) # 2 x h x w # grid = grid.reshape([2, 1, grid_size, grid_size]) grid = torch.einsum("chw,n->cnhw", grid, res) # 2 x n x h x w _, n, h, w = grid.shape pos_embed = get_2d_sincos_pos_embed_from_grid_torch(embed_dim, grid) # # (nxH*W, D/2) pos_embed = pos_embed.reshape(n, h * w, embed_dim) if cls_token: pos_embed = torch.cat( [ torch.zeros([n, 1, embed_dim], device=pos_embed.device), pos_embed, ], dim=1, ) return pos_embed def get_2d_sincos_pos_embed_from_grid_torch(embed_dim, grid): assert embed_dim % 2 == 0 # use half of dimensions to encode grid_h emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid[0]) # (H*W, D/2) emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid[1]) # (H*W, D/2) emb = torch.cat([emb_h, emb_w], dim=1) # (H*W, D) return emb def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos): """ embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) """ assert embed_dim % 2 == 0 omega = torch.arange(embed_dim // 2, device=pos.device) / embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) pos = pos.reshape(-1) # (M,) out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product emb_sin = torch.sin(out) # (M, D/2) emb_cos = torch.cos(out) # (M, D/2) emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D) return emb def get_month_encoding_table(embed_dim): """Sinusoid month encoding table, for 12 months indexed from 0-11""" assert embed_dim % 2 == 0 angles = torch.arange(0, 13) / (12 / (2 * np.pi)) sin_table = torch.sin(torch.stack([angles for _ in range(embed_dim // 2)], axis=-1)) cos_table = torch.cos(torch.stack([angles for _ in range(embed_dim // 2)], axis=-1)) month_table = torch.concatenate([sin_table[:-1], cos_table[:-1]], axis=-1) return month_table # (M, D) def adjust_learning_rate( optimizer, epoch, warmup_epochs, total_epochs, max_lr, min_lr, ): """Decay the learning rate with half-cycle cosine after warmup""" if epoch < warmup_epochs: lr = max_lr * epoch / warmup_epochs else: lr = min_lr + (max_lr - min_lr) * 0.5 * ( 1.0 + math.cos(math.pi * (epoch - warmup_epochs) / (total_epochs - warmup_epochs)) ) for group in optimizer.param_groups: group["lr"] = lr return lr # thanks to https://github.com/bwconrad/flexivit/ for this nice implementation # of the FlexiPatchEmbed module def to_2tuple(x: Any) -> Tuple: if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): return tuple(x) return tuple(itertools.repeat(x, 2)) class FlexiPatchEmbed(nn.Module): def __init__( self, patch_size: Union[int, Tuple[int, int]], in_chans: int = 3, embed_dim: int = 128, norm_layer: Optional[nn.Module] = None, bias: bool = True, patch_size_seq: Sequence[int] = (1, 2, 3, 4, 5, 6), interpolation: str = "bicubic", antialias: bool = True, ) -> None: """2D image to patch embedding w/ flexible patch sizes Extended from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/patch_embed.py#L24 by https://github.com/bwconrad/flexivit/ Args: patch_size: Base patch size. i.e the size of the parameter buffer in_chans: Number of input image channels embed_dim: Network embedding dimension size norm_layer: Optional normalization layer bias: Whether to use bias in convolution patch_size_seq: List of patch sizes to randomly sample from interpolation: Resize interpolation type antialias: Whether to apply antialiasing resizing """ super().__init__() self.patch_size = to_2tuple(patch_size) self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=bias, ) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() # Flexi specific attributes self.interpolation = interpolation self.antialias = antialias self.patch_size_seq = patch_size_seq # Pre-calculate pinvs self.pinvs = self._cache_pinvs() def _cache_pinvs(self) -> dict: """Pre-calculate all pinv matrices""" pinvs = {} for ps in self.patch_size_seq: tuple_ps = to_2tuple(ps) pinvs[tuple_ps] = self._calculate_pinv(self.patch_size, tuple_ps) return pinvs def _resize(self, x: Tensor, shape: Tuple[int, int]) -> Tensor: x_resized = F.interpolate( x[None, None, ...], shape, mode=self.interpolation, antialias=self.antialias, ) return x_resized[0, 0, ...] def _calculate_pinv(self, old_shape: Tuple[int, int], new_shape: Tuple[int, int]) -> Tensor: mat = [] for i in range(np.prod(old_shape)): basis_vec = torch.zeros(old_shape) basis_vec[np.unravel_index(i, old_shape)] = 1.0 mat.append(self._resize(basis_vec, new_shape).reshape(-1)) resize_matrix = torch.stack(mat) return torch.linalg.pinv(resize_matrix) def resize_patch_embed(self, patch_embed: Tensor, new_patch_size: Tuple[int, int]): """Resize patch_embed to target resolution via pseudo-inverse resizing""" # Return original kernel if no resize is necessary if self.patch_size == new_patch_size: return patch_embed # Calculate pseudo-inverse of resize matrix if new_patch_size not in self.pinvs: self.pinvs[new_patch_size] = self._calculate_pinv(self.patch_size, new_patch_size) pinv = self.pinvs[new_patch_size] pinv = pinv.to(patch_embed.device) def resample_patch_embed(patch_embed: Tensor): h, w = new_patch_size resampled_kernel = pinv @ patch_embed.reshape(-1) return rearrange(resampled_kernel, "(h w) -> h w", h=h, w=w) v_resample_patch_embed = vmap(vmap(resample_patch_embed, 0, 0), 1, 1) return v_resample_patch_embed(patch_embed) def forward( self, x: Tensor, patch_size: Optional[Union[int, Tuple[int, int]]] = None, ) -> Union[Tensor, Tuple[Tensor, Tuple[int, int]]]: # x has input shape [b, h, w, (t), c] batch_size = x.shape[0] has_time_dimension = False num_timesteps = 0 # ignored if has_time_dimension is False if len(x.shape) == 5: has_time_dimension = True num_timesteps = x.shape[3] x = rearrange(x, "b h w t c -> (b t) c h w") else: x = rearrange(x, "b h w c -> b c h w") if not patch_size: # During evaluation use base patch size if not specified patch_size = self.patch_size patch_size = to_2tuple(patch_size) # Resize conv weights if patch_size == self.patch_size: weight = self.proj.weight else: weight = self.resize_patch_embed(self.proj.weight, patch_size) # Apply conv with resized weights x = F.conv2d(x, weight, bias=self.proj.bias, stride=patch_size) if has_time_dimension: x = rearrange(x, "(b t) c h w -> b h w t c", b=batch_size, t=num_timesteps) else: x = rearrange(x, "b c h w -> b h w c") x = self.norm(x) return x class Attention(nn.Module): # https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py fast_attn: Final[bool] def __init__( self, dim, num_heads=8, qkv_bias=False, qk_norm=False, attn_drop=0.0, proj_drop=0.0, norm_layer=nn.LayerNorm, cross_attn: bool = False, ): super().__init__() assert dim % num_heads == 0, "dim should be divisible by num_heads" self.num_heads = num_heads self.head_dim = dim // num_heads self.scale = self.head_dim**-0.5 self.fast_attn = hasattr(torch.nn.functional, "scaled_dot_product_attention") # FIXME self.cross_attn = cross_attn self.q = nn.Linear(dim, dim, bias=qkv_bias) self.k = nn.Linear(dim, dim, bias=qkv_bias) self.v = nn.Linear(dim, dim, bias=qkv_bias) self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) def forward(self, x, y=None, attn_mask=None): B, N, C = x.shape q = self.q(x) if y is None: assert not self.cross_attn k = self.k(x) v = self.v(x) else: assert self.cross_attn k = self.k(y) v = self.v(y) q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) q, k = self.q_norm(q), self.k_norm(k) if self.fast_attn: if attn_mask is not None: attn_mask = attn_mask[:, None, None].repeat((1, self.num_heads, N, 1)) x = F.scaled_dot_product_attention( q, k, v, # a value of True indicates that the element should take part in attention attn_mask=attn_mask, dropout_p=self.attn_drop.p, ) else: if attn_mask is not None: raise NotImplementedError q = q * self.scale attn = q @ k.transpose(-2, -1) attn = attn.softmax(dim=-1) attn = self.attn_drop(attn) x = attn @ v x = x.transpose(1, 2).reshape(B, N, C) x = self.proj(x) x = self.proj_drop(x) return x class Mlp(nn.Module): """MLP as used in Vision Transformer, MLP-Mixer and related networks""" def __init__( self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.0, ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) self.act = act_layer() self.drop1 = nn.Dropout(drop) self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) self.drop2 = nn.Dropout(drop) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.drop1(x) x = self.fc2(x) x = self.drop2(x) return x class LayerScale(nn.Module): def __init__(self, dim, init_values=1e-5, inplace=False): super().__init__() self.inplace = inplace self.gamma = nn.Parameter(init_values * torch.ones(dim)) def forward(self, x): return x.mul_(self.gamma) if self.inplace else x * self.gamma def drop_path(x, drop_prob: float = 0.0, training: bool = False): if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) random_tensor.floor_() # binarize output = x.div(keep_prob) * random_tensor return output class DropPath(nn.Module): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" def __init__(self, drop_prob=None): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, x): return drop_path(x, self.drop_prob, self.training) class Block(nn.Module): def __init__( self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_norm=False, drop=0.0, attn_drop=0.0, drop_path=0.0, init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm, cross_attn: bool = False, ): super().__init__() self.norm1 = norm_layer(dim) self.attn = Attention( dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=drop, norm_layer=norm_layer, cross_attn=cross_attn, ) self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) self.mlp = Mlp( in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop, ) self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() def forward(self, x, y, attn_mask): x = x + self.drop_path(self.ls1(self.attn(self.norm1(x), y, attn_mask))) x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x)))) return x class ModuleListWithInit(nn.ModuleList): def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) class GalileoBase(nn.Module): cross_attn: bool def __init__( self, embedding_size: int = 128, depth=2, mlp_ratio=2, num_heads=8, max_sequence_length=24, base_patch_size: int = 4, use_channel_embs: bool = True, drop_path: float = 0.0, band_layout: Optional[Dict[str, Any]] = None, ): super().__init__() band_layout = band_layout or GalileoConfig().band_layout() self.space_time_groups = band_layout["space_time_groups"] self.space_groups = band_layout["space_groups"] self.time_groups = band_layout["time_groups"] self.static_groups = band_layout["static_groups"] self.base_gsd = band_layout["input_resolution_m"] self.embedding_size = embedding_size self.base_patch_size = base_patch_size self.blocks = ModuleListWithInit( [ Block( embedding_size, num_heads, mlp_ratio, qkv_bias=True, norm_layer=nn.LayerNorm, cross_attn=self.cross_attn, drop_path=drop_path, ) for _ in range(depth) ] ) self.max_sequence_length = max_sequence_length # we have 4 embeddings (pos_in_time, pos_in_space, month, channel) so each get # 0.25 of the dimension. This will change soon anyway self.pos_embed = nn.Parameter( get_1d_sincos_pos_embed_from_grid_torch( int(embedding_size * 0.25), torch.arange(max_sequence_length) ), requires_grad=False, ) month_tab = get_month_encoding_table(int(embedding_size * 0.25)) self.month_embed = nn.Embedding.from_pretrained(month_tab, freeze=True) if use_channel_embs: args = {"requires_grad": True} else: args = {"requires_grad": False} self.s_t_channel_embed = nn.Parameter( torch.zeros(len(self.space_time_groups), int(embedding_size * 0.25)), **args ) self.sp_channel_embed = nn.Parameter( torch.zeros(len(self.space_groups), int(embedding_size * 0.25)), **args ) self.t_channel_embed = nn.Parameter( torch.zeros(len(self.time_groups), int(embedding_size * 0.25)), **args ) self.st_channel_embed = nn.Parameter( torch.zeros(len(self.static_groups), int(embedding_size * 0.25)), **args ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) @classmethod def collapse_and_combine_hwtc( cls, s_t_x: torch.Tensor, sp_x: torch.Tensor, t_x: torch.Tensor, st_x: torch.Tensor, s_t_m: torch.Tensor, sp_m: torch.Tensor, t_m: torch.Tensor, st_m: torch.Tensor, ): s_t_x = rearrange(s_t_x, "b h w t c_g d -> b (h w t c_g) d") sp_x = rearrange(sp_x, "b h w c_g d -> b (h w c_g) d") t_x = rearrange(t_x, "b t c_g d -> b (t c_g) d") s_t_m = rearrange(s_t_m, "b h w t c_g-> b (h w t c_g)") sp_m = rearrange(sp_m, "b h w c_g-> b (h w c_g)") t_m = rearrange(t_m, "b t c_g -> b (t c_g)") x = torch.cat( [ s_t_x, sp_x, t_x, st_x, ], dim=1, ) m = torch.cat([s_t_m, sp_m, t_m, st_m], dim=1) return x, m @classmethod def split_and_expand_hwtc( cls, x: torch.Tensor, h: int, w: int, t: int, s_t_c_g: int, sp_c_g: int, t_c_g: int, st_c_g: int, ): n_s_t_t = h * w * t * s_t_c_g n_t_t = t * t_c_g s_t_x = rearrange(x[:, :n_s_t_t], "b (h w t c) d -> b h w t c d", h=h, w=w, t=t, c=s_t_c_g) sp_x = rearrange( x[:, n_s_t_t : -(n_t_t + st_c_g)], "b (h w c) d -> b h w c d", h=h, w=w, c=sp_c_g ) t_x = rearrange(x[:, -(n_t_t + st_c_g) : -st_c_g], "b (t c) d -> b t c d", t=t, c=t_c_g) st_x = x[:, -st_c_g:] return s_t_x, sp_x, t_x, st_x def apply_encodings(self, s_t_x, sp_x, t_x, st_x, months, patch_size, input_res): b, h, w, t, s_t_c_g, _ = s_t_x.shape sp_c_g, t_c_g = sp_x.shape[-2], t_x.shape[-2] st_c_g = st_x.shape[-2] s_t_channel = repeat(self.s_t_channel_embed, "c_g d -> b h w t c_g d", b=b, h=h, w=w, t=t) t_channel = repeat(self.t_channel_embed, "c_g d -> b t c_g d", b=b, t=t) st_channel = repeat(self.st_channel_embed, "c_g d -> b c_g d", b=b) sp_channel = repeat(self.sp_channel_embed, "c_g d -> b h w c_g d", b=b, h=h, w=w) pos_embed_s_t = repeat( self.pos_embed[:t], "t d -> b h w t c_g d", b=b, h=h, w=w, c_g=s_t_c_g ) m_embed_s_t = repeat( self.month_embed(months), "b t d -> b h w t c_g d", h=h, w=w, c_g=s_t_c_g ) pos_embed_t = repeat(self.pos_embed[:t], "t d -> b t c_g d", b=b, c_g=t_c_g) m_embed_t = repeat(self.month_embed(months), "b t d -> b t c_g d", c_g=t_c_g) t_zeros = torch.zeros(b, t, t_c_g, int(self.embedding_size * 0.25), device=t_x.device) sp_zeros = torch.zeros( b, h, w, sp_c_g, sp_channel.shape[-1] * 2, device=sp_channel.device, ) st_zeros = torch.zeros(b, st_c_g, st_channel.shape[-1] * 3, device=st_channel.device) # find the resolution that each token represents, which will be # the number of pixels in a patch * the resolution of each pixel if patch_size is None: patch_size = self.base_patch_size token_res = input_res * patch_size gsd_ratio = token_res / self.base_gsd assert h == w, "get_2d_sincos_pos_embed_with_resolution currently requires that h==w" spatial_embed = get_2d_sincos_pos_embed_with_resolution( int(self.embedding_size * 0.25), h, torch.ones(b).to(s_t_x.device) * gsd_ratio, device=s_t_x.device, ) spatial_embed = rearrange(spatial_embed, "b (h w) d -> b h w d", h=h, w=w) spatial_embed_s_t = repeat( spatial_embed, "b h w d -> b h w t c_g d", h=h, w=w, t=t, c_g=s_t_c_g ) spatial_embed_s = repeat(spatial_embed, "b h w d -> b h w c_g d", h=h, w=w, c_g=sp_c_g) s_t_embed = torch.cat([s_t_channel, pos_embed_s_t, m_embed_s_t, spatial_embed_s_t], dim=-1) sp_embed = torch.cat([sp_channel, sp_zeros, spatial_embed_s], dim=-1) t_embed = torch.cat([t_channel, pos_embed_t, m_embed_t, t_zeros], dim=-1) st_embed = torch.cat([st_channel, st_zeros], dim=-1) return s_t_x + s_t_embed, sp_x + sp_embed, t_x + t_embed, st_x + st_embed class Encoder(GalileoBase): cross_attn = False def __init__( self, max_patch_size: int = 8, embedding_size: int = 128, depth=2, mlp_ratio=2, num_heads=8, max_sequence_length=24, freeze_projections: bool = False, drop_path: float = 0.0, band_layout: Optional[Dict[str, Any]] = None, ): super().__init__( embedding_size, depth, mlp_ratio, num_heads, max_sequence_length, max_patch_size, use_channel_embs=True, drop_path=drop_path, band_layout=band_layout, ) self.space_time_embed = nn.ModuleDict( { group_name: FlexiPatchEmbed( in_chans=len(group), embed_dim=embedding_size, patch_size=max_patch_size ) for group_name, group in self.space_time_groups.items() } ) self.space_embed = nn.ModuleDict( { group_name: FlexiPatchEmbed( in_chans=len(group), embed_dim=embedding_size, patch_size=max_patch_size ) for group_name, group in self.space_groups.items() } ) self.time_embed = nn.ModuleDict( { group_name: nn.Linear(in_features=len(group), out_features=embedding_size) for group_name, group in self.time_groups.items() } ) self.static_embed = nn.ModuleDict( { group_name: nn.Linear(in_features=len(group), out_features=embedding_size) for group_name, group in self.static_groups.items() } ) if freeze_projections: self.space_time_embed.requires_grad_(False) self.space_embed.requires_grad_(False) self.time_embed.requires_grad_(False) self.static_embed.requires_grad_(False) self.norm = nn.LayerNorm(embedding_size) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) def apply_linear_projection( self, s_t_x: torch.Tensor, sp_x: torch.Tensor, t_x: torch.Tensor, st_x: torch.Tensor, s_t_m: torch.Tensor, sp_m: torch.Tensor, t_m: torch.Tensor, st_m: torch.Tensor, patch_size: int, ): """ Given a [B, H, W, (T), C] inputs, returns a [B, H, W, (T), C_G, D] output. We assume that the spatial masks are consistent for the given patch size, so that if patch_size == 2 then one possible mask would be [0, 0, 1, 1] [0, 0, 1, 1] [1, 1, 0, 0] [1, 1, 0, 0] for the H, W dimensions """ b, h, w, t, _ = s_t_x.shape new_h, new_w = h // patch_size, w // patch_size s_t_l, sp_l, t_l, st_l, s_t_m_l, sp_m_l, t_m_l, st_m_l = [], [], [], [], [], [], [], [] for idx, (channel_group, channel_idxs) in enumerate(self.space_time_groups.items()): s_t_m_l.append(s_t_m[:, 0::patch_size, 0::patch_size, :, idx]) if s_t_m_l[-1].min() == 0: s_t_l.append( self.space_time_embed[channel_group]( s_t_x[:, :, :, :, channel_idxs], patch_size=patch_size ) ) else: s_t_l.append( torch.zeros( b, new_h, new_w, t, self.embedding_size, dtype=s_t_x.dtype, device=s_t_x.device, ) ) for idx, (channel_group, channel_idxs) in enumerate(self.space_groups.items()): sp_m_l.append(sp_m[:, 0::patch_size, 0::patch_size, idx]) if sp_m_l[-1].min() == 0: sp_l.append( self.space_embed[channel_group]( sp_x[:, :, :, channel_idxs], patch_size=patch_size ) ) else: sp_l.append( torch.zeros( b, new_h, new_w, self.embedding_size, dtype=sp_x.dtype, device=sp_x.device, ) ) for idx, (channel_group, channel_idxs) in enumerate(self.time_groups.items()): t_m_l.append(t_m[:, :, idx]) if t_m_l[-1].min() == 0: t_l.append(self.time_embed[channel_group](t_x[:, :, channel_idxs])) else: t_l.append( torch.zeros(b, t, self.embedding_size, dtype=t_x.dtype, device=t_x.device) ) for idx, (channel_group, channel_idxs) in enumerate(self.static_groups.items()): st_m_l.append(st_m[:, idx]) if st_m_l[-1].min() == 0: st_l.append(self.static_embed[channel_group](st_x[:, channel_idxs])) else: st_l.append( torch.zeros(b, self.embedding_size, dtype=st_x.dtype, device=st_x.device) ) return ( torch.stack(s_t_l, dim=-2), torch.stack(sp_l, dim=-2), torch.stack(t_l, dim=-2), torch.stack(st_l, dim=-2), torch.stack(s_t_m_l, dim=-1), torch.stack(sp_m_l, dim=-1), torch.stack(t_m_l, dim=-1), torch.stack(st_m_l, dim=-1), ) @staticmethod def remove_masked_tokens(x, mask): org_mask_dtype = mask.dtype mask = mask.bool() # https://stackoverflow.com/a/68621610/2332296 # move all non-masked values to the front of their rows sorted_mask, indices = torch.sort((~mask).int(), dim=1, descending=True, stable=True) x = x.gather(1, indices[:, :, None].expand_as(x)) # set masked values to 0 (not really necessary since we'll ignore them anyway) x = x * sorted_mask.unsqueeze(-1) # cut off to the length of the longest sequence max_length = sorted_mask.sum(-1).max() x = x[:, :max_length] updated_mask = 1 - sorted_mask[:, :max_length] return x, indices, updated_mask.to(dtype=org_mask_dtype) @staticmethod def add_removed_tokens(x, indices, mask): masked_tokens = repeat( torch.zeros_like(x[0, 0, :]), "d -> b t d", b=x.shape[0], t=indices.shape[1] ) full_mask = torch.cat( ( mask, torch.ones( (x.shape[0], indices.shape[1] - x.shape[1]), device=x.device, dtype=mask.dtype ), ), dim=-1, ) # can't set value on leaf variable out = masked_tokens.clone() # put tokens in full masked tensor (at the first N positions in every row) out[~full_mask.bool()] = x[~mask.bool()] # then move them to their original positions out = out.scatter(1, indices[:, :, None].expand_as(out), out) full_mask = full_mask.scatter(1, indices.expand_as(full_mask), full_mask) return out, full_mask def apply_attn( self, s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, months, patch_size, input_res, exit_after, token_exit_cfg, ): if token_exit_cfg: exit_s_t, exit_sp, exit_t, exit_st = self.create_token_exit_ids( s_t_x, sp_x, t_x, st_x, token_exit_cfg ) exit_ids_seq, _ = self.collapse_and_combine_hwtc( exit_s_t, exit_sp, exit_t, exit_st, s_t_m, sp_m, t_m, st_m ) # exited_tokens starts as linear projections! exited_tokens, _ = self.collapse_and_combine_hwtc( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m ) else: exit_ids_seq = None exited_tokens = None _, h, w, t, s_t_c_g, _ = s_t_x.shape sp_c_g, t_c_g, st_c_g = sp_x.shape[3], t_x.shape[-2], st_x.shape[-2] s_t_x, sp_x, t_x, st_x = self.apply_encodings( s_t_x, sp_x, t_x, st_x, months, patch_size, input_res ) x, m = self.collapse_and_combine_hwtc(s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m) # we only care about the values >= 1 for this mask, since 2 just tells the decoder # to decode those tokens. From the perspective of the encoder, 1 and 2 are equivalent # since they both represent masked values new_m = m >= 1 x, indices, new_m = self.remove_masked_tokens(x, new_m) # new_m is shape (bsz, seq_len) if exit_ids_seq is not None: exit_ids_seq, _, _ = self.remove_masked_tokens(exit_ids_seq, m >= 1) # still linear projections exited_tokens, _, _ = self.remove_masked_tokens(exited_tokens, m >= 1) for i_blk, blk in enumerate(self.blocks): if (exit_after is not None) and ((i_blk + 1) > exit_after): # if exit_after is N, then we exit after the Nth layer # if exit_after is 0, then all layers are skipped break # skip the 0th block since this is just the linear # projection if (exit_ids_seq is not None) and (i_blk > 0): assert exited_tokens is not None # half depth exited_tokens = torch.where( condition=(exit_ids_seq == i_blk), input=x.detach(), other=exited_tokens.detach(), ) # we take the inverse of the mask because a value # of True indicates the value *should* take part in # attention x = blk(x=x, y=None, attn_mask=~new_m.bool()) if exit_ids_seq is not None: assert exited_tokens is not None # full depth # IMPORTANT: write this to x x = torch.where( condition=(exit_ids_seq == (i_blk + 1)), # 2 for full depth input=x.detach(), other=exited_tokens.detach(), ) # we don't care about the mask returned by add_removed_tokens, since we will # just use the original, unclipped mask here x, _ = self.add_removed_tokens(x, indices, new_m) return ( *self.split_and_expand_hwtc(x, h, w, t, s_t_c_g, sp_c_g, t_c_g, st_c_g), s_t_m, sp_m, t_m, st_m, ) @classmethod def average_tokens(cls, s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m): x, m = cls.collapse_and_combine_hwtc(s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m) x, _, m = cls.remove_masked_tokens(x, m) x_for_mean = x * (1 - m.unsqueeze(-1)) return x_for_mean.sum(dim=1) / torch.sum(1 - m, -1, keepdim=True) @classmethod def apply_mask_and_average_tokens_per_patch( cls, s_t_x: torch.Tensor, sp_x: torch.Tensor, t_x: torch.Tensor, st_x: torch.Tensor, s_t_m: torch.Tensor, sp_m: torch.Tensor, t_m: torch.Tensor, st_m: torch.Tensor, ): s_t_x = rearrange(s_t_x, "b t_h t_w t c_g d -> b (t_h t_w) (t c_g) d") sp_x = rearrange(sp_x, "b t_h t_w c_g d -> b (t_h t_w) c_g d") # repeat time tokens over space t_x = repeat( rearrange(t_x, "b t c_g d -> b (t c_g) d"), "b n d -> b s n d", s=sp_x.shape[1] ) st_x = repeat(st_x, "b c_g d -> b s c_g d", s=sp_x.shape[1]) s_t_m = rearrange(s_t_m, "b t_h t_w t c_g-> b (t_h t_w) (t c_g)") sp_m = rearrange(sp_m, "b t_h t_w c_g-> b (t_h t_w) c_g") t_m = repeat(rearrange(t_m, "b t c_g -> b (t c_g)"), "b n -> b s n", s=sp_x.shape[1]) st_m = repeat(st_m, "b c_g -> b s c_g", s=sp_x.shape[1]) x = torch.cat([s_t_x, sp_x, t_x, st_x], dim=2) # B, S, N, D m = torch.cat([s_t_m, sp_m, t_m, st_m], dim=2) # B, S, N x_for_mean = x * (1 - m.unsqueeze(-1)) return x_for_mean.sum(dim=2) / torch.sum(1 - m, -1, keepdim=True) def create_token_exit_ids(self, s_t_x, sp_x, t_x, st_x, token_exit_cfg): exit_s_t = torch.zeros_like(s_t_x) exit_sp = torch.zeros_like(sp_x) exit_t = torch.zeros_like(t_x) exit_st = torch.zeros_like(st_x) for idx, (key, _) in enumerate(self.space_time_groups.items()): exit_s_t[:, :, :, :, idx, :] = token_exit_cfg[key] for idx, (key, _) in enumerate(self.space_groups.items()): exit_sp[:, :, :, idx, :] = token_exit_cfg[key] for idx, (key, _) in enumerate(self.time_groups.items()): exit_t[:, :, idx, :] = token_exit_cfg[key] for idx, (key, _) in enumerate(self.static_groups.items()): exit_st[:, idx, :] = token_exit_cfg[key] return exit_s_t, exit_sp, exit_t, exit_st def forward( self, s_t_x: torch.Tensor, sp_x: torch.Tensor, t_x: torch.Tensor, st_x: torch.Tensor, s_t_m: torch.Tensor, sp_m: torch.Tensor, t_m: torch.Tensor, st_m: torch.Tensor, months: torch.Tensor, patch_size: int, input_resolution_m: Optional[int] = None, exit_after: Optional[int] = None, token_exit_cfg: Optional[Dict] = None, add_layernorm_on_exit: bool = True, ): if input_resolution_m is None: input_resolution_m = self.base_gsd ( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, ) = self.apply_linear_projection( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, patch_size ) if (exit_after is None) or (exit_after > 0): s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m = self.apply_attn( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, months, patch_size, input_resolution_m, exit_after=exit_after, token_exit_cfg=token_exit_cfg, ) if add_layernorm_on_exit: s_t_x = self.norm(s_t_x) sp_x = self.norm(sp_x) t_x = self.norm(t_x) st_x = self.norm(st_x) return ( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, months, ) @classmethod def load_from_folder( cls, folder: Path, device: torch.device, config_filename: str = "config.json", encoder_filename: str = "encoder.pt", ): if not (folder / config_filename).exists(): all_files_in_folder = [f.name for f in folder.glob("*")] raise ValueError( f"Expected {config_filename} in {folder}, found {all_files_in_folder}" ) if not (folder / encoder_filename).exists(): all_files_in_folder = [f.name for f in folder.glob("*")] raise ValueError( f"Expected {encoder_filename} in {folder}, found {all_files_in_folder}" ) with (folder / config_filename).open("r") as f: config = json.load(f) model_config = config["model"] encoder_config = model_config["encoder"] encoder = cls(**encoder_config) state_dict = torch.load(folder / encoder_filename, map_location=device) for key in list(state_dict.keys()): # this cleans the state dict, which occasionally had an extra # ".backbone" included in the key names state_dict[key.replace(".backbone", "")] = state_dict.pop(key) encoder.load_state_dict(state_dict) return encoder logger = logging.get_logger(__name__) def _default_pretraining_normalizing_dict() -> Dict[str, Dict[str, List[float]]]: return { "13": { "mean": [ -11.728724389184965, -18.85558188024017, 1395.3408730676722, 1338.4026921784578, 1343.09883810357, 1543.8607982512297, 2186.2022069512263, 2525.0932853316694, 2410.3377187373408, 2750.2854646886753, 2234.911100061487, 1474.5311266077113, 0.2892116502999044, ], "std": [ 4.887145774840316, 5.730270320384293, 917.7041440370853, 913.2988423581528, 1092.678723527555, 1047.2206083460424, 1048.0101611156767, 1143.6903026819996, 1098.979177731649, 1204.472755085893, 1145.9774063078878, 980.2429840007796, 0.2720939024500081, ], }, "16": { "mean": [ 673.0152819503361, 5.930092668915115, 0.10470439140978786, 0.23965913270066183, 0.08158044385860364, 0.04246976254259546, 0.11304392863520317, 0.17329647890362473, 0.0698981691616277, 0.12130267132802142, 0.04671318615236216, 10.973119802517362, 1.0927069179958768, 1.6991394232855903, 0.03720594618055555, 1.3671352688259548, ], "std": [ 983.0697298296237, 8.167406789813247, 0.18771647977504985, 0.2368313455675914, 0.08024268534756586, 0.04045374496146404, 0.11350342472061795, 0.1279898111718168, 0.12042341550438586, 0.13602408145504347, 0.043971116096060345, 31.255340146970997, 10.395974878206689, 12.92380617159917, 1.9285254295940466, 11.612179775408928, ], }, "6": { "mean": [ 271.5674963541667, 0.08554303677156568, 657.3181260091111, 692.1291795806885, 562.781331880633, 1.5647115934036673, ], "std": [ 79.80828940314429, 0.11669547098151486, 704.0008695557707, 925.0116126406431, 453.2434022278578, 7.513020170832818, ], }, "18": { "mean": [ 188.20315880851746, 0.2804946561574936, 0.11371652073860168, 0.058778801321983334, 0.10474256777763366, 0.2396918488264084, 0.08152248692512512, 0.04248040814399719, 0.11303179881572724, 0.17326324067115784, 0.06998309404850006, 0.12122812910079957, 0.04671641788482666, 10.98456594619751, 1.0968475807189941, 1.6947754135131836, 0.03320046615600586, 1.3602827312469483, ], "std": [ 1154.5919128300602, 0.5276998078079327, 0.7021637331734328, 0.36528892213195063, 0.17470213191865785, 0.20411195416718833, 0.0660782470089761, 0.03380702424871257, 0.09809195568521663, 0.11292471052124119, 0.09720748930233268, 0.12912217763726777, 0.0399973913151906, 23.725471823867462, 5.715238079725388, 9.030481416228302, 0.9950220242487364, 7.754429123862099, ], }, } class GalileoConfig(PreTrainedConfig): model_type = "galileo" def __init__( self, hidden_size: int = 128, num_hidden_layers: int = 4, num_attention_heads: int = 8, mlp_ratio: float = 4.0, max_sequence_length: int = 24, max_patch_size: int = 8, freeze_projections: bool = False, drop_path: float = 0.1, default_patch_size: int = 8, default_month: int = 6, global_pool: bool = True, input_resolution_m: int = 10, s1_bands: Optional[List[str]] = None, s2_bands: Optional[List[str]] = None, era5_bands: Optional[List[str]] = None, tc_bands: Optional[List[str]] = None, viirs_bands: Optional[List[str]] = None, srtm_bands: Optional[List[str]] = None, dw_bands: Optional[List[str]] = None, wc_bands: Optional[List[str]] = None, landscan_bands: Optional[List[str]] = None, location_bands: Optional[List[str]] = None, space_time_band_groups: Optional[Dict[str, List[str]]] = None, time_band_groups: Optional[Dict[str, List[str]]] = None, space_band_groups: Optional[Dict[str, List[str]]] = None, pretraining_normalizing_dict: Optional[Dict[str, Dict[str, List[float]]]] = None, **kwargs, ): super().__init__(**kwargs) self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.mlp_ratio = mlp_ratio self.max_sequence_length = max_sequence_length self.max_patch_size = max_patch_size self.freeze_projections = freeze_projections self.drop_path = drop_path self.default_patch_size = default_patch_size self.default_month = default_month self.global_pool = global_pool self.input_resolution_m = input_resolution_m self.s1_bands = s1_bands if s1_bands is not None else ["VV", "VH"] self.s2_bands = s2_bands if s2_bands is not None else [ "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12", ] self.era5_bands = era5_bands if era5_bands is not None else [ "temperature_2m", "total_precipitation_sum", ] self.tc_bands = tc_bands if tc_bands is not None else ["def", "soil", "aet"] self.viirs_bands = viirs_bands if viirs_bands is not None else ["avg_rad"] self.srtm_bands = srtm_bands if srtm_bands is not None else ["elevation", "slope"] self.dw_bands = dw_bands if dw_bands is not None else [ "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops", "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice", ] self.wc_bands = wc_bands if wc_bands is not None else [ "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation", ] self.landscan_bands = landscan_bands if landscan_bands is not None else ["b1"] self.location_bands = location_bands if location_bands is not None else ["x", "y", "z"] self.space_time_band_groups = space_time_band_groups if space_time_band_groups is not None else { "S1": ["VV", "VH"], "S2_RGB": ["B2", "B3", "B4"], "S2_Red_Edge": ["B5", "B6", "B7"], "S2_NIR_10m": ["B8"], "S2_NIR_20m": ["B8A"], "S2_SWIR": ["B11", "B12"], "NDVI": ["NDVI"], } self.time_band_groups = time_band_groups if time_band_groups is not None else { "ERA5": ["temperature_2m", "total_precipitation_sum"], "TC": ["def", "soil", "aet"], "VIIRS": ["avg_rad"], } self.space_band_groups = space_band_groups if space_band_groups is not None else { "SRTM": ["elevation", "slope"], "DW": [ "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops", "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice", ], "WC": [ "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation", ], } self.pretraining_normalizing_dict = ( pretraining_normalizing_dict if pretraining_normalizing_dict is not None else _default_pretraining_normalizing_dict() ) def band_layout(self) -> Dict[str, Any]: space_time_bands = self.s1_bands + self.s2_bands + ["NDVI"] time_bands = self.era5_bands + self.tc_bands + self.viirs_bands space_bands = self.srtm_bands + self.dw_bands + self.wc_bands static_dw_bands = [f"{band}_static" for band in self.dw_bands] static_wc_bands = [f"{band}_static" for band in self.wc_bands] static_bands = self.landscan_bands + self.location_bands + static_dw_bands + static_wc_bands static_band_groups = { "LS": self.landscan_bands, "location": self.location_bands, "DW_static": static_dw_bands, "WC_static": static_wc_bands, } return { "s1_bands": self.s1_bands, "s2_bands": self.s2_bands, "era5_bands": self.era5_bands, "tc_bands": self.tc_bands, "viirs_bands": self.viirs_bands, "srtm_bands": self.srtm_bands, "dw_bands": self.dw_bands, "wc_bands": self.wc_bands, "landscan_bands": self.landscan_bands, "location_bands": self.location_bands, "space_time_bands": space_time_bands, "time_bands": time_bands, "space_bands": space_bands, "static_bands": static_bands, "space_time_groups": _band_group_indices(space_time_bands, self.space_time_band_groups), "time_groups": _band_group_indices(time_bands, self.time_band_groups), "space_groups": _band_group_indices(space_bands, self.space_band_groups), "static_groups": _band_group_indices(static_bands, static_band_groups), "input_resolution_m": self.input_resolution_m, } class GalileoPreTrainedModel(PreTrainedModel): config_class = GalileoConfig base_model_prefix = "galileo" main_input_name = "space_time_x" input_modalities = ("image",) supports_gradient_checkpointing = True _no_split_modules = ["Block"] _supports_sdpa = False class GalileoEncoderModel(GalileoPreTrainedModel): def __init__(self, config: GalileoConfig, add_pooling_layer: bool = True): super().__init__(config) self.encoder = Encoder( max_patch_size=config.max_patch_size, embedding_size=config.hidden_size, depth=config.num_hidden_layers, mlp_ratio=config.mlp_ratio, num_heads=config.num_attention_heads, max_sequence_length=config.max_sequence_length, freeze_projections=config.freeze_projections, drop_path=config.drop_path, band_layout=config.band_layout(), ) self.add_pooling_layer = add_pooling_layer and config.global_pool self.post_init() def forward( self, space_time_x=None, space_x=None, time_x=None, static_x=None, space_time_mask=None, space_mask=None, time_mask=None, static_mask=None, months=None, patch_size=None, input_resolution_m=None, return_dict=None, **kwargs: Unpack[TransformersKwargs], ): if space_time_x is None: raise ValueError("You must specify `space_time_x`") if return_dict is None: return_dict = self.config.use_return_dict patch_size = patch_size if patch_size is not None else self.config.default_patch_size input_resolution_m = ( input_resolution_m if input_resolution_m is not None else self.config.input_resolution_m ) outputs = self.encoder( space_time_x, space_x, time_x, static_x, space_time_mask, space_mask, time_mask, static_mask, months, patch_size=patch_size, input_resolution_m=input_resolution_m, ) s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m, _ = outputs last_hidden_state, mask = Encoder.collapse_and_combine_hwtc( s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m ) last_hidden_state, _, _ = Encoder.remove_masked_tokens(last_hidden_state, mask) pooler_output = None if self.add_pooling_layer: pooler_output = Encoder.average_tokens(s_t_x, sp_x, t_x, st_x, s_t_m, sp_m, t_m, st_m) if not return_dict: return (last_hidden_state, pooler_output) return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooler_output, ) __all__ = ["Encoder", "GalileoConfig", "GalileoEncoderModel", "GalileoPreTrainedModel"]