Intro to GeoBenchV2 datasets

This notebook is intended to illustrate the common features that the GeoBenchV2 datasets have and which hopefully allow users to more easily configure a range of different experiments they are interested in. We will choose the BenV2 dataset for demonstration purposes in this tutorial

from pathlib import Path

import torch

from geobench_v2.datamodules import GeoBenchBENV2DataModule
from geobench_v2.datasets import GeoBenchBENV2

GeoBenchV2 provides LightningDataModules for all datasets. If you are not familiar with Lightning yet, you are encouraged to follow their documentation and tutorials. However, in summary, a LightningDataModule organizes the train, validation and test dataset with common functionality under one module. This ensures consistency across datasets and helps with reproducability.

The datasets are implemented as PyTorch datasets, which can be wrapped with a PyTorch DataLoader to yield machine learning batches for training. The datamodule organizes all this into a common structure.

DATA_ROOT = Path("../../data/benv2")
band_order = GeoBenchBENV2.band_default_order

datamodule = GeoBenchBENV2DataModule(
    img_size=120,
    batch_size=4,
    num_workers=4,
    root=DATA_ROOT,
    band_order=band_order,
    data_normalizer=torch.nn.Identity(),  # we do custom normalization in the tutorial
)
datamodule.setup("fit")
datamodule.setup("test")

Using provided pre-initialized normalizer instance: Identity
Using provided pre-initialized normalizer instance: Identity
Using provided pre-initialized normalizer instance: Identity

You can access the dataloaders for each split as follows:

train_dataloader = datamodule.train_dataloader()
val_dataloader = datamodule.val_dataloader()
test_dataloader = datamodule.test_dataloader()

EO-specific features

As you probably know, Earth Observational (EO) vision data is different than other common Computer Vision data.

Band Selection

We aim to make the configuration of band selection easy by supporting different options:

• specify a dictionary that includes the desired sensor modality names as keys and the desired bands as a list

• the order of band names also dictates the order in which the bands are returned in the data loader

• additionally, one can also specify fill values (floats) that will yield a channel with just this value

These band selection features are intended to aid experiments that for example look into the importance of band ordering or channel importance. We will highlight the three cases.

dataset = GeoBenchBENV2(
    root=DATA_ROOT,
    split="train",
    data_normalizer=torch.nn.Identity(),  # we do custom normalization in the tutorial
)

Using provided pre-initialized normalizer instance: Identity

band_order_one = {"s1": ["VV", "VH"], "s2": ["B04", "B03", "B02"]}

dataset_one = GeoBenchBENV2(
    root=DATA_ROOT,
    split="train",
    band_order=band_order_one,
    data_normalizer=torch.nn.Identity(),
)

band_order_two = {"s1": ["VV", "VH"], "s2": ["B02", "B03", "B04"]}

dataset_two = GeoBenchBENV2(
    root=DATA_ROOT,
    split="train",
    band_order=band_order_two,
    data_normalizer=torch.nn.Identity(),
)

band_order_three = {"s1": ["VV", 0.0, "VH"], "s2": ["B04", "B03", 1.5, "B02"]}

dataset_three = GeoBenchBENV2(
    root=DATA_ROOT,
    split="train",
    band_order=band_order_three,
    data_normalizer=torch.nn.Identity(),
)

Using provided pre-initialized normalizer instance: Identity
Using provided pre-initialized normalizer instance: Identity
Using provided pre-initialized normalizer instance: Identity

sample_one, sample_two, sample_three = (
    dataset_one[0],
    dataset_two[0],
    dataset_three[0],
)

Normalization

Input normalization plays an important role in deep learning training dynamics. Since EO data does usually not have a constrained input range like standard camera RGB imagery (0-255 pixel values), different normalization schemes exist. To easily configure this, each dataset class has a data_normalizer argument which is intended to be a PyTorch module that normalizes each sample according to the band statistics. We provide two types of normalizations, the classic Z-Score normalization ZScoreNormalizer, as well as a normalization scheme popularized by SatMAE and hence termed SatMAENormalizer. A small deep-dive and demonstration of these normalization schemes is provided in this tutorial. We have precomputed normalization and target statistics according to each scheme that are available alongside the datasets on Huggingface and are automatically downloaded with the dataset.

Augmentations

It is common to utilize data augmentations to increase the diversity of the training data and improve optimization objectives. In the Lightning framework such augmentations are recommended to be efficiently applied on an entire batch of data, even on GPU in the on_after_batch_transfer method. They can be configured separately for training and evaluation through the train_augmentations (applied to training set) and eval_augmentations (applied to both validation and test set) arguments. However, these augmentations will only be applied when also conducting model training within the Lightning framework. Since users might want to use the “raw” PyTorch datasets within their own setup or further evaluation schemes, we have consciously separated normalization and augmentations, even though they are often treated as one. For example if you are coming from TorchGeo, note that in their DataModules, both normalization and augmentations are applied under the augmentations parameter in on_after_batch_transfer function.

Data Visualization

Visualizations are an incredibly useful but sometimes under utilized tool for debugging and model training inspection, as well as more fundamentally becoming familiar with the data to understand potential pitfalls that could explain model behavior. For these reasons, each DataModule has a visualize_batch method to which you can either pass a batch of data, or simply ask the method to visualize samples from the train, validation, or test split. For both classification, and segmentation task, there are also utilization functions that visualize the target statistics. The use of these is demonstrated in each of the dataset notebooks.