Open in Studio Open in Colab 
[ ]:

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

Command-Line Interface#

Written by: Adam J. Stewart and Caleb Robinson

TorchGeo provides a command-line interface based on LightningCLI that allows users to combine our data modules and trainers from the comfort of the command line. This no-code solution can be attractive for both beginners and experts, as it offers flexibility and reproducibility. In this tutorial, we demonstrate some of the features of this interface.

Setup#

First, we install TorchGeo. In addition to the Python library, this also installs a torchgeo executable.

[ ]:

!uv pip install torchgeo

Subcommands#

The torchgeo command has a number of subcommands that can be run. The --help flag can be used to list them.

[ ]:

!torchgeo --help

Trainer#

Below, we run --help on the fit subcommand to see what options are available to us. fit is used to train and validate a model, and we can customize many aspects of the training process.

[ ]:

!torchgeo fit --help

Model#

We must first select an nn.Module model architecture to train and a lightning.pytorch.LightningModule trainer to train it. We will experiment with the ClassificationTask trainer and see what options we can customize. Any of TorchGeo’s builtin trainers, or trainers written by the user, can be used in this way.

[ ]:

!torchgeo fit --model.help ClassificationTask

Data#

We must also select a Dataset we would like to train on and a lightning.pytorch.LightningDataModule we can use to access the train/val/test split and any augmentations to apply to the data. Similarly, we use the --help flag to see what options are available for the EuroSAT100 dataset.

[ ]:

!torchgeo fit --data.help EuroSAT100DataModule

Config#

Now that we have seen all important configuration options, we can put them together in a YAML file. LightingCLI supports YAML, JSON, and command-line configuration. While we will write this file using Python in this tutorial, normally this file would be written in your favorite text editor.

[ ]:

import os
import tempfile

root = os.path.join(tempfile.gettempdir(), 'cli')
config = f"""
trainer:
  max_epochs: 2
  default_root_dir: '{root}'
  callbacks:
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        monitor: val_loss
        mode: min
        save_top_k: 1
        save_last: true
        filename: 'eurosat-{{epoch:02d}}-{{val_loss:.4f}}'
    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
      init_args:
        logging_interval: epoch

model:
  class_path: ClassificationTask
  init_args:
    model: 'resnet18'
    weights: 'ResNet18_Weights.SENTINEL2_ALL_MOCO'
    in_channels: 13
    num_classes: 10

data:
  class_path: EuroSAT100DataModule
  init_args:
    batch_size: 8
  dict_kwargs:
    root: '{root}'
    download: true

optimizer:
  class_path: torch.optim.AdamW
  init_args:
    lr: 1e-3
    weight_decay: 0.01

lr_scheduler:
  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
  init_args:
    T_max: 2
    eta_min: 1e-6
"""
os.makedirs(root, exist_ok=True)
with open(os.path.join(root, 'config.yaml'), 'w') as f:
    f.write(config)

This YAML file has five potential sections:

  • trainer: Arguments to pass to the Trainer, including callbacks

  • model: Arguments to pass to the task

  • data: Arguments to pass to the data module

  • optimizer (optional): Custom optimizer configuration

  • lr_scheduler (optional): Learning rate scheduler configuration

The class_path gives the class to instantiate and init_args lists its arguments. For data modules, dict_kwargs can pass additional keyword arguments to the underlying dataset.

Optimizer and Learning Rate Scheduler#

By default, TorchGeo trainers use AdamW with a ReduceLROnPlateau scheduler. However, you can customize these via the optimizer and lr_scheduler sections in your config.

In our example, we use:

  • AdamW: Adam with decoupled weight decay regularization, which often leads to better generalization than standard Adam

  • CosineAnnealingLR: Gradually reduces the learning rate following a cosine curve from the initial lr down to eta_min. This smooth decay helps the model converge to a better minimum. The T_max parameter should match your max_epochs

Other popular scheduler choices include:

  • torch.optim.lr_scheduler.StepLR - Decay LR by a factor every N epochs

  • torch.optim.lr_scheduler.OneCycleLR - Super-convergence with warmup and annealing

  • torch.optim.lr_scheduler.LinearLR - Linear warmup (often combined with cosine via SequentialLR)

Callbacks#

Lightning callbacks allow you to hook into the training loop at various points. In this example config, we use two callbacks:

ModelCheckpoint#

ModelCheckpoint automatically saves model weights during training. Key options include:

  • monitor: The metric to track (e.g., val_loss, val_Accuracy)

  • mode: Whether to minimize (min) or maximize (max) the monitored metric

  • save_top_k: Number of best checkpoints to keep (set to -1 to save all)

  • save_last: Also save the most recent checkpoint regardless of performance

  • filename: Template for checkpoint filenames with metric placeholders

LearningRateMonitor#

LearningRateMonitor logs the current learning rate to your logger which is useful to make sure your learning rate schedule is working as expected.

Other Useful Callbacks#

Lightning provides many other callbacks that you can add to your config:

  • EarlyStopping: Stop training when a metric stops improving

    - class_path: lightning.pytorch.callbacks.EarlyStopping
  init_args:
    monitor: val_loss
    patience: 10
    mode: min

  • StochasticWeightAveraging: Improve generalization by averaging weights

    - class_path: lightning.pytorch.callbacks.StochasticWeightAveraging
  init_args:
    swa_lrs: 1e-4

  • GradientAccumulationScheduler: Simulate larger batch sizes on limited GPU memory

    - class_path: lightning.pytorch.callbacks.GradientAccumulationScheduler
  init_args:
    scheduling: {0: 1, 4: 2, 8: 4}  # Accumulate more as training progresses

  • DeviceStatsMonitor: Log GPU memory usage and utilization

    - class_path: lightning.pytorch.callbacks.DeviceStatsMonitor

See the Lightning Callbacks documentation for more options.

Training#

We can now train our model like so.

[ ]:

!torchgeo fit --config {root}/config.yaml

Validation#

Now that we have a trained model, we can evaluate performance on the validation set. Note that we need to explicitly pass in the location of the checkpoint from the previous run.

[ ]:

import glob

checkpoint = glob.glob(
    os.path.join(root, 'lightning_logs', 'version_0', 'checkpoints', '*.ckpt')
)[0]

!torchgeo validate --config {root}/config.yaml --ckpt_path {checkpoint}

Testing#

After finishing our hyperparameter tuning, we can calculate and report the final test performance.

[ ]:

!torchgeo test --config {root}/config.yaml --ckpt_path {checkpoint}

Additional Reading#

Lightning CLI has many more features that are worth learning. You can learn more by reading the following set of tutorials: