---

name: datamodules-and-tasks description: "Use for TorchGeo Lightning DataModules, CLI/config runs, task classes, losses, metrics, prediction/plot hooks, and training workflow wiring." metadata: disco-role: operating disable-model-invocation: true license: MIT

TorchGeo datamodules and tasks

Use this sub-skill when wiring TorchGeo into Lightning training, validation, testing, prediction, CLI configs, or task modules.

DataModule operating facts

• BaseDataModule wraps a dataset class, batch_size, num_workers, split datasets, split dataloaders, and Kornia augmentation hooks.
• prepare_data() may download data when download=True; it must not create split state. setup(stage) creates datasets/samplers and is called on every process.
• GeoDataModule is for GeoDataset classes and sets collate_fn = stack_samples. It manages RandomPatchSampler, GriddedPatchSampler, and batch samplers for split-specific patch extraction.
• NonGeoDataModule is for integer-indexed datasets and ordinary PyTorch dataloader shuffle/split semantics.
• on_after_batch_transfer() applies the split-specific augmentation after the batch reaches the device. This is where Kornia GPU augmentations belong.

Task classes

TorchGeo task modules are Lightning modules for common geospatial ML problems:

• Classification: timm-backed image classification with weights, in_channels, task type, class/label counts, losses, and optional frozen backbone.
• SemanticSegmentation: SMP or TorchGeo FCN segmentation with model, backbone, weights, in_channels, multiclass/multilabel/binary settings, losses, and freeze options.
• Other task modules cover change detection, detection, instance segmentation, regression, temporal/spatiotemporal tasks, MAE, MoCo, BYOL, SimCLR, and IO benchmark workflows.

Before changing a task, inspect its corresponding tests/tasks/test_<task>.py for constructor matrix, batch schema, metric expectations, and trainer smoke coverage.

CLI/config workflow

• The torchgeo console script is declared as torchgeo.main:main.
• Prefer JSONArgParse/Lightning config files for reproducible training. Explicitly name the datamodule, task/model, trainer, optimizer/scheduler parameters, and paths.
• For smoke checks, use small fake fixtures and max_epochs=1, small batch sizes, and no checkpoint downloads.

Batch and mask handling

• Images are typically tensors shaped (B, C, H, W); some spatiotemporal paths use (B, T, C, H, W) and may flatten or route to spatiotemporal tasks.
• Segmentation masks may require dtype conversion for Kornia and restoration before loss/metrics.
• Detection and instance segmentation batches can contain variable-length boxes/masks; use the datamodule utility collate function used by the tests.

Read next

• reference for task-specific examples and verification candidates.
• Root troubleshooting for Lightning, Kornia, and optional dependency failures.