# Copyright (c) TorchGeo Contributors. All rights reserved.
# Licensed under the MIT License.

Lightning Trainers

Written by: Caleb Robinson

In this tutorial, we demonstrate TorchGeo trainers to train and test a model. We will use the EuroSAT dataset throughout this tutorial. Specifically, a subset containing only 100 images. We will train models to predict land cover classes.

It’s recommended to run this notebook on Google Colab if you don’t have your own GPU. Click the “Open in Colab” button above to get started.

Setup

First, we install TorchGeo and TensorBoard.

%pip install torchgeo tensorboard

Imports

Next, we import TorchGeo and any other libraries we need.

%matplotlib inline
%load_ext tensorboard

import os
import tempfile

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger

from torchgeo.datamodules import EuroSAT100DataModule
from torchgeo.models import ResNet18_Weights
from torchgeo.trainers import ClassificationTask

Lightning modules

Our trainers use Lightning to organize both the training code, and the dataloader setup code. This makes it easy to create and share reproducible experiments and results.

First we’ll create a EuroSAT100DataModule object which is simply a wrapper around the EuroSAT100 dataset. This object 1.) ensures that the data is downloaded, 2.) sets up PyTorch DataLoader objects for the train, validation, and test splits, and 3.) ensures that data from the same region is not shared between the training and validation sets so that you can properly evaluate the generalization performance of your model.

The following variables can be modified to control training.

batch_size = 10
num_workers = 2
max_epochs = 50
fast_dev_run = False

root = os.path.join(tempfile.gettempdir(), 'eurosat100')
datamodule = EuroSAT100DataModule(
    root=root, batch_size=batch_size, num_workers=num_workers, download=True
)

torchgeo.trainers provides specialized task classes that simplify training workflows for common geospatial tasks. Depending on your objective, you can select the appropriate trainer class, such as ClassificationTask for classification, SemanticSegmentationTask for semantic segmentation, or other task-specific trainers. Check the trainers documentation for more information.

Given that EuroSAT is a classification dataset, we can use a ClassificationTask object that holds the model and optimizer as well as the training logic.

Here, we create a ClassificationTask object that holds the model object, optimizer object, and training logic. We will use a ResNet-18 model that has been pre-trained on Sentinel-2 imagery.

task = ClassificationTask(
    loss='ce',
    model='resnet18',
    weights=ResNet18_Weights.SENTINEL2_ALL_MOCO,
    in_channels=13,
    num_classes=10,
    lr=0.1,
    patience=5,
)

Training

Now that we have the Lightning modules set up, we can use a Lightning Trainer to run the training and evaluation loops. There are many useful pieces of configuration that can be set in the Trainer – below we set up model checkpointing based on the validation loss, early stopping based on the validation loss, and a TensorBoard based logger. All checkpoints and logs will be stored in the default_root_dir directory. We encourage you to see the Lightning docs for other options that can be set here, e.g. CSV logging, automatically selecting your optimizer’s learning rate, and easy multi-GPU training.

default_root_dir = os.path.join(tempfile.gettempdir(), 'experiments')
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss', dirpath=default_root_dir, save_top_k=1, save_last=True
)
early_stopping_callback = EarlyStopping(monitor='val_loss', min_delta=0.0, patience=10)
logger = TensorBoardLogger(save_dir=default_root_dir, name='tutorial_logs')

For tutorial purposes we deliberately lower the maximum number of training epochs.

trainer = Trainer(
    callbacks=[checkpoint_callback, early_stopping_callback],
    fast_dev_run=fast_dev_run,
    log_every_n_steps=1,
    logger=logger,
    min_epochs=1,
    max_epochs=max_epochs,
)

When we first call .fit(...) the dataset will be downloaded and checksummed (if it hasn’t already). After this, the training process will kick off, and results will be saved so that TensorBoard can read them.

trainer.fit(model=task, datamodule=datamodule)

We launch TensorBoard to visualize various performance metrics across training and validation epochs. We can see that our model is just starting to converge, and would probably benefit from additional training time and a lower initial learning rate.

%tensorboard --logdir "$default_root_dir"

Finally, after the model has been trained, we can easily evaluate it on the test set. If you train several models with different hyperparameters, you can select the one with the best validation performance using ckpt_path='best'.

trainer.test(model=task, datamodule=datamodule)