fix: OlmoEarth backbone forward pass to match olmoearth_pretrain API - #1202
Open
whatnick wants to merge 2 commits into
Open
fix: OlmoEarth backbone forward pass to match olmoearth_pretrain API#1202whatnick wants to merge 2 commits into
whatnick wants to merge 2 commits into
Conversation
Add integration for OlmoEarth (Allen AI) foundation models as registered backbones in the TerraTorch backbone registry. - Register 11 model variants: v1 (nano/tiny/base/large), v1.1 (nano/tiny/base), v1.2 (nano/tiny/small/base) - OlmoEarthBackbone wrapper adapts the encoder's token output to spatial feature maps (B, D, H', W') for TerraTorch decoders - Added as optional dependency: pip install terratorch[olmoearth] - Tests with skipif guard for environments without olmoearth-pretrain References: - GitHub: https://github.com/allenai/olmoearth_pretrain - Paper: https://arxiv.org/abs/2511.13655 - Issue: torchgeo#1201 Signed-off-by: Tisham Dhar <tisham.dhar@csiro.au>
- Use correct modality name 'sentinel2_l2a' (NamedTuple field name) - Pass raw pixels in [B, H, W, T, C] format (encoder does patchification) - Use pixel-resolution mask with correct num_band_sets dimension - Use long dtype timestamps with [day, month, year] format - Read embed_dim from encoder.embedding_size instead of hardcoded config - Extract spatial features from tokens_and_masks (not project_aggregated) - Fix Nano embed_dim from 64 to 128 in config and tests Signed-off-by: Tisham Dhar <tisham.dhar@csiro.au>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the OlmoEarth backbone integration to work correctly with the
olmoearth-pretrainpackage API.Changes
sentinel2_l2a(matchingMaskedOlmoEarthSampleNamedTuple fields) instead of the invalidsentinel2[B, H, W, T, C]format — the encoder handles patchification internally[B, H, W, T, num_band_sets]with correct band set count (3 for S2 L2A)[B, T, 3]long tensors with[day, month, year](not float zeros)encoder.embedding_sizeat runtime instead of relying on potentially stale config valuestokens_and_masks(averaged over T and band_sets), not the globalproject_aggregatedvectorTesting
All 8 tests in
tests/test_olmoearth.pypass (~2 min on CPU, no GPU required):