Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
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.
First, we install TorchGeo and TensorBoard.
%pip install torchgeo tensorboard
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
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
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, )
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.
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.