Contribute a model (Stage 2)#

This guide covers everything a pull request needs to add a new frozen pretrained model to torchgeo-bench. Before continuing, complete Evaluate your own model (Stage 1) which demonstrates how to implement your new model to work with the torchgeo-bench pipeline and verify your model produces sensible results on the applicable datasets.

Prerequisites#

The setup steps are the same as Stage 1:

$ git clone https://github.com/torchgeo/torchgeo-bench.git
$ cd torchgeo-bench
$ conda activate torchgeo-bench
$ uv sync --extra dev

Then fork the repository on GitHub and create a feature branch:

$ git remote add fork https://github.com/<your-username>/torchgeo-bench.git
$ git checkout -b add-<model-name>

For the model implementation and YAML config steps, follow Evaluate your own model (Stage 1) (Sections Implement your model and Create a Hydra config). The remainder of this page covers integration, tests, and the PR submission.

Integrate into the package#

Once your model class is working locally, move it into torchgeo-bench and export it so the Hydra registry can resolve _target_.

1. Place the module under src/torchgeo_bench/models/:

$ mv new_model.py src/torchgeo_bench/models/new_model.py

2. Export the class from src/torchgeo_bench/models/__init__.py:

# src/torchgeo_bench/models/__init__.py
from .new_model import NewModel

__all__: list[str] = [
    # ... existing entries (keep alphabetical) ...
    "NewModel",
]

3. Update the Hydra config _target_ to the package path:

# src/torchgeo_bench/conf/model/new_model.yaml
_target_: torchgeo_bench.models.NewModel
name: new_model
pretrained: true

4. Declare optional dependencies in pyproject.toml if your model requires packages beyond [project.dependencies]:

[project.optional-dependencies]
newmodel = ["newpackage>=1.0"]

Install the extra locally to confirm it resolves:

$ uv sync --extra newmodel

Note

If your model requires an optional extra, document it clearly in your model’s __init__ docstring and in the PR description. The test suite must still import the class without the extra installed (guard weight loading with a try/except ImportError only around the optional import, not the class definition).

Weights#

Pretrained weights must be publicly accessible without authentication. HuggingFace Hub is the preferred host.
The model must load after a fresh pip install 'torchgeo-bench[newextra]' with no manual file placement. Use huggingface_hub.hf_hub_download or an equivalent auto-download call inside your __init__.
The weights URL must appear in the PR description so reviewers can verify provenance.

Write tests#

Create tests/test_<model>.py. Every added code path must be covered.

Fast tests (run in CI) — no network I/O:

import torch
import pytest
from torchgeo_bench.datasets.base import BandSpec
from torchgeo_bench.models.new_model import NewModel


def _bands(n: int = 3) -> list[BandSpec]:
    return [
        BandSpec(sensor="s2", name=f"b{i}", source_name=f"B{i}",
                 mean=500.0, std=100.0, min=0.0, max=10000.0)
        for i in range(n)
    ]


def test_new_model_output_shape():
    """Model returns (B, K) with random weights."""
    model = NewModel(bands=_bands(), pretrained=False)
    x = torch.randn(2, 3, 64, 64)
    with torch.no_grad():
        out = model.forward_patch_features(x)
    assert out.ndim == 2
    assert out.shape[0] == 2


def test_new_model_num_channels():
    """`num_channels` matches the input BandSpec list length."""
    model = NewModel(bands=_bands(5), pretrained=False)
    assert model.num_channels == 5

Weight-download tests (slow, run locally before PR) — mark with @pytest.mark.slow:

@pytest.mark.slow
def test_new_model_pretrained_loads():
    """Pretrained weights download and load without error."""
    model = NewModel(bands=_bands(), pretrained=True)
    x = torch.randn(1, 3, 64, 64)
    with torch.no_grad():
        out = model.forward_patch_features(x)
    assert out.ndim == 2

Run the fast tests before opening the PR:

$ pytest --no-cov tests/test_new_model.py

Slow tests must pass locally but are excluded from the default CI run (pytest without -m slow skips them automatically):

$ pytest --no-cov -m slow tests/test_new_model.py

Submit results#

Run the full benchmark on all datasets applicable to your model’s sensor coverage and write the results to results/contributed/<model_name>.csv:

$ torchgeo-bench run model=new_model \
    dataset.names=[m-eurosat,m-so2sat,m-bigearthnet,m-brick-kiln,m-forestnet,m-pv4ger] \
    output=results/contributed/new_model.csv

For V2 datasets:

$ torchgeo-bench run model=new_model \
    dataset.names=[benv2,treesatai,so2sat,forestnet] \
    output=results/contributed/new_model.csv resume=true

The CSV schema is identical to results/all_results.csv — see Results format for the full column reference.

Lint and full test suite#

Before opening the PR, apply auto-fixes and verify the full test suite passes:

$ ruff check . --fix && ruff format .
$ pytest --no-cov

Open the PR#

When all checklist items are satisfied, open a pull request against main using the “Add model” template:

.github/PULL_REQUEST_TEMPLATE/add_model.md

All checklist items must be checked before requesting a review. The template prompts you for:

A model summary table (name, pretraining data, sensor coverage, weights URL, and paper/project page if available).
Confirmation that each technical requirement is satisfied (class exported, config present, weights public, tests written and passing, results submitted, lint clean).