diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py index 5ea16a4..136886f 100644 --- a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py +++ b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py @@ -840,7 +840,7 @@ def __init__( self, embedding_size: int, supported_modalities: list[ModalitySpec], - max_sequence_length: int, + max_sequence_length: int | None = None, learnable_channel_embeddings: bool = True, random_channel_embeddings: bool = False, tokenization_config: TokenizationConfig | None = None, @@ -853,7 +853,8 @@ def __init__( embedding_size: Size of token embeddings supported_modalities: Which modalities from Modality this model instantiation supports - max_sequence_length: Maximum sequence length + max_sequence_length: Size of the checkpoint-compatible temporal + encoding table. Longer sequences are computed on-the-fly. learnable_channel_embeddings: Whether to use learnable channel embeddings random_channel_embeddings: Initialize channel embeddings randomly (zeros if False) tokenization_config: Optional config for custom band groupings @@ -870,6 +871,8 @@ def __init__( f"position_encoding must be one of {PositionEncoding.values()}, " f"got {position_encoding}" ) + if max_sequence_length is None: + max_sequence_length = 12 self.embedding_size = embedding_size self.supported_modalities = supported_modalities self.supported_modality_names = [ @@ -878,15 +881,13 @@ def __init__( self.tokenization_config = tokenization_config or TokenizationConfig() self.position_encoding = position_encoding self.embedding_size = embedding_size - self.max_sequence_length = ( - max_sequence_length # This max sequence length is a time dim thing - ) # TODO: we need to be able to calculate the size of the param based on what types of embeddings it will get # we have 4 embeddings types (pos_in_time, pos_in_space, month, channel) so each get # 0.25 of the dimension self.embedding_dim_per_embedding_type = int(embedding_size * 0.25) - # Position encodings for time dimension initialized to 1D sinusoidal encodings + # Retain the frozen table for exact compatibility with existing + # checkpoints. Positions beyond it are computed on-the-fly. self.pos_embed = nn.Parameter( get_1d_sincos_pos_encoding( torch.arange(max_sequence_length), @@ -1032,10 +1033,15 @@ def _apply_encodings_per_modality( # Slot-index temporal encoding (additive). Skipped when 3D RoPE # handles temporal position rotationally inside attention. if not PositionEncoding.is_3d_rope(self.position_encoding): - time_embed = repeat( - self.pos_embed[:t], f"t d -> {ein_string}", **ein_dict - ) - modality_embed[..., n : n * 2] += time_embed.to(device) + if t <= self.pos_embed.shape[0]: + pos_embed = self.pos_embed[:t].to(device) + else: + pos_embed = get_1d_sincos_pos_encoding( + torch.arange(t, device=device), + self.embedding_dim_per_embedding_type, + ) + time_embed = repeat(pos_embed, f"t d -> {ein_string}", **ein_dict) + modality_embed[..., n : n * 2] += time_embed # Month encodings stay additive in all modes (calendar/seasonal # signal is orthogonal to slot-index). diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py index 9f903a7..c35c040 100644 --- a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py +++ b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/constants.py @@ -4,7 +4,7 @@ """ from dataclasses import dataclass -from enum import Enum +from enum import StrEnum # The highest resolution that we are working at. # Everything else is a factor (which is a power of 2) coarser than this resolution. @@ -81,7 +81,7 @@ def get_expected_image_size(self, modality_resolution_factor: int) -> int: return IMAGE_TILE_SIZE // (self.resolution_factor // modality_resolution_factor) -class TimeSpan(str, Enum): +class TimeSpan(StrEnum): """Enum to distinguish data that is valid for different time ranges.""" # Only one data point (not time series). diff --git a/tests/test_flexi_vit.py b/tests/test_flexi_vit.py index faf7d25..0c75c7a 100644 --- a/tests/test_flexi_vit.py +++ b/tests/test_flexi_vit.py @@ -1,7 +1,58 @@ -from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import Encoder +import torch + +from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.encodings import ( + get_1d_sincos_pos_encoding, +) +from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import ( + CompositeEncodings, + Encoder, +) from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality +def test_composite_encodings_compute_temporal_positions_on_the_fly() -> None: + """CompositeEncodings supports sequences longer than deprecated config value.""" + encodings = CompositeEncodings( + embedding_size=16, + supported_modalities=[Modality.SENTINEL2_L2A], + max_sequence_length=3, + ) + assert encodings.pos_embed.shape == (3, 4) + + tokens = torch.zeros((1, 2, 2, 5, Modality.SENTINEL2_L2A.num_band_sets, 16)) + timestamps = torch.tensor( + [[[1, 0, 2020], [2, 1, 2020], [3, 2, 2020], [4, 3, 2020], [5, 4, 2020]]] + ) + + output = encodings( + {Modality.SENTINEL2_L2A.name: tokens}, + timestamps=timestamps, + patch_size=4, + ) + + encoded = output[Modality.SENTINEL2_L2A.name] + assert encoded.shape == tokens.shape + assert torch.isfinite(encoded).all() + + expected_time = get_1d_sincos_pos_encoding(torch.arange(5), 4) + actual_time = encoded[0, 0, 0, :, 0, 4:8] + assert torch.allclose(actual_time, expected_time, atol=1e-5) + + +def test_composite_encodings_loads_temporal_position_table() -> None: + """Strict state dict loading restores the frozen temporal table.""" + encodings = CompositeEncodings( + embedding_size=16, + supported_modalities=[Modality.SENTINEL2_L2A], + max_sequence_length=3, + ) + state_dict = encodings.state_dict() + state_dict["pos_embed"] = torch.zeros((3, 4)) + + encodings.load_state_dict(state_dict, strict=True) + assert torch.equal(encodings.pos_embed, torch.zeros((3, 4))) + + def test_band_dropout_disabled_by_default() -> None: """Test that Encoder leaves band dropout disabled at construction.""" encoder = Encoder(