Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Lightning Trainers

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

import torch
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
)

Next, 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. 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.

accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'
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.00, 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(
    accelerator=accelerator,
    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.

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