This post is part of my portfolio of personal projects resulting from my transition from generalist data science (2013-2025) to geospatial data science, GIS, and GeoAI (2025+).
Using a segmentation neural network to detect objects in satellite raster data is now a fairly common GeoAI use case, and can be done in a broad-brush sense with no-code or low-code tools. Finer-grained control of the model, data, and processing is often still desirable for real end-to-end analyses, however. The TorchGeo tool builds upon PyTorch and PyTorch Lightning to enable end-to-end GeoAI for a wide range of use cases. We show how plain PyTorch code is augmented by Lightning’s data and model modules, then in turn by TorchGeo’s geospatial functionality. This lets us run end-to-end building detection as an example, and control the details.
GeoAI and the cloud are the next frontiers of GIS and geospatial data science, allowing powerful analyses not previously possible.
Here we train and use a model to detect buildings in the Inria data supplied with TorchGeo. While the data and functions can largely be used as black boxes, it is informative to show how TorchGeo is building on several layers of capabilities to make end-to-end GeoAI as tractable as it is here.
We cover
- PyTorch neural network with DataLoader, training loop, etc.
- PyTorch Lightning with Data and Model modules
- TorchGeo instantiations of these modules for Inria data
- TorchGeo calling of these for end-to-end data preparation, model training, and inference
The code to run this end-to-end is available in this project’s GitHub repository.
Setup
End-to-end GeoAI in exploratory as opposed to production mode is a workflow that benefits from being able to combine code, data visualization, and commentary inline. In other words, it is helpful to use a modern form of literate programming commonly known as a Jupyter notebook.
We use the VS Code IDE which allows Python and Jupyter to be run within a virtual environment, along with the functionalities of debugging, inline class documentation, AI coding assistance, etc.
AI assistance in particular can help rapidly prototype code that is approximately syntactically correct. But it still in general needs to be refined to be 100% what we want, especially in terms of doing the correct analysis. Similarly to a built-in LLM model generating code, an AI agent could potentially be prompted to do the whole analysis, but the same issues of code and analysis correctness would be encountered, with errors potentially extending end-to-end rather than over one notebook cell.
A Python virtual environment .venv is created, within which we can install TorchGeo using their preferred pip install torchgeo route. This in turn brings in the needed dependencies of PyTorch, PyTorch Lightning, and other libraries. VS Code’s Python and Jupyter extensions enable notebook use.
The author’s Mac laptop has an Apple M3 Max processor with GPU, and since our purposes here are exploratory and pedagogical it was not necessary to access cloud-scale GPU resources to run things. Amazon EC2 was accessed, but it quickly became apparent that modern GPUs such as A100 or H100 were in short supply. Providers such as Lambda, CoreWeave, or DigitalOcean would be alternative options if we wanted to pursue the cloud, along with geospatial tiling tools to run things at scale.
Similarly, because the data are not large, we download it to local disk rather than accessing remote storage such as an Amazon S3 cloud bucket.
Geospatial machine learning
Most machine learning is designed to work with either tabular data, textual data, or images. For geospatial, two main modifications are needed:
- Spatial awareness of models w.r.t. the images or other information
- Handling color images with more channels than the usual red, green, and blue (RGB)
This means that most non-geo tutorials won’t work as-is with geo data, and need modifying at the code level. Changes include handling images within the data loader, the model architecture, the training loop, and inference.
A major benefit of TorchGeo is that it streamlines these modifications, while remaining open source with the underlying code visible. This enables geo data to be used end-to-end in the PyTorch ecosystem without having to write the code from scratch, but in a manner that can be understood as opposed to a black box.
Building up to Torchgeo
It is instructive to build up to what TorchGeo is doing by looking at a typical model in PyTorch, followed by PyTorch Lightning, and then TorchGeo.
PyTorch
The basics of end-to-end usage of a model are well illustrated by the tutorial quick-start in the PyTorch GitHub repository. Let’s run through the code, with an eye on our upcoming geo end-to-end analysis. (Quoting the code here is allowed by the repository’s LICENSE file.) This model runs a small neural network on the classic non-geo dataset Fashion MNIST, essentially a “hello world” end-to-end.
Assuming everything is installed (TorchGeo will bring in these dependencies), import the needed modules from PyTorch and TorchVision:
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensorDownload the data. Fashion MNIST is a built-in TorchVision dataset, hence datasets:
training_data = datasets.FashionMNIST(root="data", train=True, download=True, transform=ToTensor())
test_data = datasets.FashionMNIST(root="data", train=False, download=True, transform=ToTensor())This gives us the data as a Dataset object, in tensor format suitable for a model. It also has metadata about whether it is used for training, and thus has labels, or testing.
When a model is trained, it is not given all the data at once, so Dataset objects are wrapped with a DataLoader. This iteratively goes through the Dataset to pass it to the model, and handles other needed operations such as randomizing the order of the rows.
batch_size = 64
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)
for X, y in test_dataloader:
print(f"Shape of X [N, C, H, W]: {X.shape}")
print(f"Shape of y: {y.shape} {y.dtype}")
breakThe for loop shows us the data shape so we know that it is what it should be: N images, C channels, each image of height H and width W. This is important for geo data since there will often be more channels than the usual 3 in RGB color images.
Now we define the model. When using pure PyTorch like here, this is done by creating a model class which subclasses the nn.module class. We can therefore see the full model architecture and change any details if we need to. The device line puts the model onto a GPU if one is available.
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28*28, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, 10)
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
model = NeuralNetwork().to(device)
print(model)The network is defined by __init__, and how the data passes through is defined by forward.
The model is then part of the training loop, which passes the data forward through the model, and the model weights are updated by the backward pass shown below. The model is performing well when the weights are such that the error on its predictions of the labels in the testing set, the loss, is minimized.
The model is put through the training loop a number of times, or epochs, and so we have the train() and test() functions, and a loop on epochs.
Training:
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
# Backpropagation
loss.backward()
optimizer.step()
optimizer.zero_grad()
if batch % 100 == 0:
loss, current = loss.item(), (batch + 1) * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")Testing:
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")Epochs:
epochs = 5
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn)
print("Done!")When the model has been trained we want to save it:
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")Often when training large models, intermediate steps called checkpoints are saved during training so that if something goes wrong the run does not have to be restarted from the beginning.
Now that there is a trained model, we want to perform the step that many tutorials and articles do not cover: apply the model to new data!
PyTorch can do this in a simple way by using eval
model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))
classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
x = x.to(device)
pred = model(x)
predicted, actual = classes[pred[0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')Here the new data still has labels, but is not required to.
In a production system, the trained model can be deployed on a server, and be run on new incoming data, i.e., inference.
PyTorch Lightning
From the above we see that we need to
- Obtain the data
- Prepare the data
- Define a data loader
- Define the model
- Run the training loop
- Evaluate the model
- Run inference on new data
PyTorch Lightning is designed to modularize and streamline this process so that the code remains manageable when the functionality of the basic model above is extended to more complex data and models.
Geo data will have changes from non-geo data in all of the above steps, so Lightning is a big plus for handling all of the necessary processing.
Lightning’s modularization is in the form of
- Dataset: Type (e.g., raster), file locations, plots, etc.
- Data module: Data preprocessing and loading into the model
- Lightning module: Which training task to run: classification, segmentation, etc.
- Trainer: Model settings for training
- Trainer.fit: Run training (validation, testing also available)
- Trainer.predict: Run inference on new data
Schematically, this looks like
dataset = DatasetName(files, ...)
datamodule = DatasetNameDataModule(dataset, ...)
task = NeuralNetworkModel(model, settings, hyperparameters, ...)
trainer = Trainer(model settings)
trainer.fit(task, datamodule, ...)
trainer.predict(task, datamodule, ...)This allows different combinations of datasets and models to be run without having to rewrite a lot of code.
TorchGeo
TorchGeo in turn uses the PyTorch Lightning modular approach and supplies ready-made dataset, data loader, and model modules for many geospatial datasets. It also includes various library functions to ease the geospatial processing components of this, such as combining images of parts of an area (common because images are large) into a larger image of the whole area.
So in TorchGeo we have
- Datasets: PyTorch dataset, geo or non-geo, raster or vector, etc., from curated or customizable sources
- Data modules: Download data and create a DataLoader for a specific geo dataset with appropriate splits such as train/test in non-overlapping geographic areas
- Models: ML models that have appropriate architecture and good performance with geo data
- Trainers: Run a given ML model to perform a given task, such as classification, segmentation, etc., followed by inference
There are many further features like custom losses, samplers that take into account bounding boxes, and transforms to add geo featurizations to data such as vegetation indices. The geo processing functions can then do things like combining Datasets.
It is this as a whole which allows us to say let’s run the U-Net segmentation model on the Inria Aerial Image Labeling dataset to detect buildings, and then apply it to some unseen data, without it being a major project. But we retain access to the details given by the full code, rather than their being hidden behind a low- or no-code interface. Such interfaces are great if they do exactly what you want, but they become much less useful if they don’t, and often completely opaque if there is an error.
Running the building detection model
Now we have an overview of how everything works, let’s run the building detection model. For the full code, see this project’s GitHub repository.
Define the model task we want to run
task = SemanticSegmentationTask( # SemanticSegmentation in 0.10+
model='unet',
backbone='resnet50',
weights=True, # Use existing pretrained ImageNet weights
in_channels=3, # RGB
task='binary', # 2 classes building and not-building
loss='bce', # TorchGeo 0.9+ adds pos_weight and Dice loss
freeze_backbone=True, # Freeze encoder and train decoder + head
lr=0.001 # Learning rate
)Choose the data we want to use, the ready-defined data module Inria
datamodule = InriaAerialImageLabelingDataModule(
batch_size=1, # Large batches go out of memory on M3 laptop
patch_size=1024,
num_workers=2,
root=data_dir
)Here, batch_size is how many training examples are passed to the model each time around the training loop. We used a batch size of 1 due to memory limitations on the laptop. The data module breaks up each of Inria’s 5000×5000 pixel TIFF images into smaller patches of, e.g., 64×64 pixels that can be passed through the model.
Log and visualize the training using TensorBoard, helpful for, e.g., checking that the validation loss doesn’t start becoming worse than the training loss (overfitting):
%load_ext tensorboard
logger = TensorBoardLogger(save_dir=experiments_dir, name='inria_logs')
%tensorboard --logdir "$experiments_dir"Set model and early stopping checkpoints. These provide safety against having to rerun the model from scratch, and the training going on for too long, respectively. Early stopping is mostly illustrative as we do not run for many epochs on the laptop, although it did get invoked in the final results we show below.
checkpoint_callback = ModelCheckpoint(
monitor='val_loss',
dirpath=experiments_dir,
save_top_k=1, # Save best model
save_last=True,
mode='min' # Minimize the monitored quantity (val_loss)
)
early_stopping_callback = EarlyStopping(
monitor='val_loss',
min_delta=0.0,
patience=10, # Stop if not significant change after this many epochs
mode='min',
strict=True,
check_finite=True
)Set up the model trainer
trainer = Trainer(
accelerator='auto', # Can see gpu or mps on Mac
callbacks=[checkpoint_callback, early_stopping_callback],
default_root_dir=experiments_dir,
fast_dev_run=fast_dev_run,
max_epochs=20,
limit_train_batches=frac_subsample,
limit_val_batches=frac_subsample,
limit_test_batches=frac_subsample,
limit_predict_batches=frac_subsample,
log_every_n_steps=1,
logger=logger
)where frac_subsample < 1 can be used for faster sanity-check runs, max_epochs can be set, and accelerator automatically puts the model on the GPU.
Training, model validation, and predictions (inference) on unseen data can then be run in one line each:
trainer.fit(model=task, datamodule=datamodule)results_val = trainer.validate(model=task, datamodule=datamodule)preds = trainer.predict(model=task, datamodule=datamodule)Viewing the Results
In principle that gives us all our results. Of course, we want to view them, and this takes some more coding. Writing ready-made user classes for such visualization is less convenient than preprocessing or training because the exact form of predictions from the model depend on the data, the model, and all the settings.
Nevertheless, TorchGeo still reduces the coding to some extent by supplying generic functionality such as a plot method in the Inria data module to show samples of images.
The Inria data consists of the training set of 155 images:
| Area | No. Images | Indices | Filename (.tif) |
| Austin, TX | 31 | 0-30 | austin |
| Chicago, IL | 31 | 31-61 | chicago |
| Kitsap county, WA | 31 | 62-92 | kitsap |
| West Tyrol, Austria | 31 | 93-123 | tyrol-w |
| Vienna, Austria | 31 | 124-154 | vienna |
the validation set of 25 images:
| Area | No. Images | Indices | Filename (.tif) |
| Austin, TX | 5 | 0-4 | austin |
| Chicago, IL | 5 | 5-9 | chicago |
| Kitsap county, WA | 5 | 10-14 | kitsap |
| West Tyrol, Austria | 5 | 15-19 | tyrol-w |
| Vienna, Austria | 5 | 20-24 | vienna |
and the testing set of 180 images:
| Area | No. Images | Indices | Filename (.tif) |
| Bellingham, WA | 36 | 0-35 | bellingham |
| Bloomington, IN | 36 | 36-71 | bloomington |
| Innsbruck, Austria | 36 | 72-107 | innsbruck |
| San Francisco, CA | 36 | 108-143 | sfo |
| East Tyrol, Austria | 36 | 144-179 | tyrol-e |
The testing set areas are not in the training or validation sets.
Note that the indices come in the order of the image filenames, which differ from the areas themselves. It always pays to check the data!
For more information about the Inria dataset, see the original paper.
Input data images can be viewed using the plot method, where 0 is an index as in the tables above:
idx_tr = 0
sample = datamodule.train_dataset[idx_tr]
datamodule.plot(sample, suptitle=f"Train dataset sample index {idx_tr} of 0-154")
Left = RGB satellite image. Right = ground truth pixels representing buildings.
To view the model predictions, we need more code than plot because we want to see all the predictions for the area. This means recombining the individual patches that were sent through the model. Our IDE’s AI assistance helps here by rapidly generating prototype code from an appropriately phrased comment that acts as a prompt. We then perform manual modification and verification to ensure it is fully correct.
Create a combined image
N = 0 # Index of image in preds to combine
def combine_patches(patches, grid_size=(np_dim, np_dim), patch_size=(patch_size, patch_size)):
combined_image = np.zeros((grid_size[0] * patch_size[0], grid_size[1] * patch_size[1]))
for i in range(grid_size[0]):
for j in range(grid_size[1]):
combined_image[i * patch_size[0]:(i + 1) * patch_size[0],
j * patch_size[1]:(j + 1) * patch_size[1]] = patches[i * grid_size[1] + j].squeeze().cpu().numpy()
return combined_image
combined_image = combine_patches(preds[N])and plot it alongside the corresponding original
sample = datamodule.predict_dataset[N]
image = sample['image'].squeeze().cpu().numpy().astype(np.uint8)
fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))
axs[0].imshow(image.transpose(1, 2, 0)) # image shape is (3, 5000, 5000), need to transpose to (5000, 5000, 3) for plt.imshow
axs[0].set_title('Original Image')
axs[0].axis('off')
axs[1].imshow(combined_image, cmap='gray')
axs[1].set_title('Combined Predictions')
axs[1].axis('off')
plt.show()giving, e.g.,
We now have our result: the trained model, and its predictions on unseen data in an area not seen during training.
The result can then be further analyzed and used depending upon what the overarching business (or other) goal is.
How good are the results?
At first glance the right-hand panel from this and other test set images look promising. We can see that the buildings are found, and they resemble the originals. TensorBoard also shows that metrics such as the loss are ok, with the validation loss following the training loss so we are not overfitting. Various other details throughout the end-to-end analysis not all reported here check out as well.
However, it is also apparent from the image, in a way that it would not be from just the numerical metrics, that the model has not only found the buildings, but also other similarly colored areas such as roads, the beach, and the barren area to the bottom right. This is seen consistently over multiple runs. Is this because we are using three channel RGB data so roads look like buildings?
Imperfect results such as this are where the ability to drill down into the details of what we are doing becomes crucial.
By viewing the settings of the data module ( InriaAerialImageLabelingDataModule), we see that the default patch size into which the image is broken up to pass through the network is 64×64 pixels. This sounds good, but for the specific task of semantic segmentation into buildings it is too small because the different shapes of the roads, i.e., lines compared to the rectangles of the buildings, are obscured. So the model is left with the RGB colors and “detecting” the roads as buildings is the best it can do.
Increasing the patch size to 1024×1024, for which my laptop has enough memory, and running the training for more epochs (the early stopping callback stopped at 16), improves the results so that now the detections are now buildings and not roads:
Note: (1) The predictions image is 4096×4096 pixels versus the 5000×5000 original. This could be resolved with further handling of the patches. (2) The plot resembles the one on TorchGeo’s GitHub page, but they used different settings, e.g., batch size 64.
This is similarly true for other images and the other areas in the Inria test set. It’s not perfect, but it is much better.
We now have our building detector with promising results. The buildings are mostly detected, and there are not many false positives. The main limitation looks to be the fidelity of the building shapes.
Improving the results further
Improving the results even more from the above is not urgent here, because our main purpose has been to show how TorchGeo is built up in turn from PyTorch and PyTorch Lightning, followed by correct end-to-end usage on real data.
Nevertheless, because we have ready access to the full details of what we ran, and the ability to customize it, we have a set of avenues that could be pursued.
A selection of possible improvements are:
- Larger model backbone: We use the
resnet50backbone for the model, which is fairly substantial. With more compute resources (memory, GPUs, cloud parallelization), the largerresnet152may benefit us. - Larger run: Running for more epochs, larger batch sizes, etc., may improve things, especially with a larger model. A quantified model hyperparameter sweep would likely be beneficial.
- More input data: Adding channels that further differentiate buildings from non-buildings that have similar colors in RGB would help if the channels have signal, perhaps satellite infrared data, building heights, or NDVI vegetation index. Infrared may not be great on its own because buildings and roads are often both concrete. Building heights are often similar to trees, but they are not vegetation. So it’s hard to be sure. Other data sources would have to be calibrated, projected, etc., to match Inria.
- Improved loss metrics: Here we run on TorchGeo 0.7.1, which supports binary cross-entropy (BCE) loss, ok for simple yes-no segmentation on buildings. Adding class weighting, because most of the image pixels are not buildings, and other metrics such as Dice that focus on Intersection-over-Union of the pixels in the buildings class would improve this. TorchGeo 0.10.0 supports these, but at the time of writing there was a bug that broke image plotting. Similarly, using 0.7.1’s
JaccardIoU option also failed. Running a bunch of versions between 0.7 and 0.10 to find some optimal bug-free combination would take more time than we want to spend here currently when a future fix is expected. - Model fine-tuning: the encoder, decoder, and decision head parts of U-Net can be tuned individually in TorchGeo with the other parts of the model frozen. This is a form of model fine-tuning. We already ran the above with the encoder frozen.
- Overlays: Showing the predictions on top of the images and highlighting pixels that should have been building but are not, or vice versa, would highlight parts of an image where the model is doing well or not so well. Similarly, viewing all 180 prediction images and deriving statistics for the different Inria areas would add insight.
- Output building vectors: Many applications of building detection would want the model outputs to be vectors of building outlines rather than pixel rasters. Geospatial rasters can readily be converted to vectors with some appropriate data processing code.
- Other details: The prediction images are 4096×4096 pixels, cutting off some of the 5000×5000 originals due to the 1024×1024 non-overlapping patches. Adding a patch stride or other handling would resolve this.
- Go modern: The ultimate real-world answer for best detection of buildings is, instead of using the decade old Inria data and a basic U-Net model, use modern embeddings data and state-of-the-art geo foundation models now available in the cloud. Examples of both are options in TorchGeo. Here, we wanted to show building up to TorchGeo, and hence better understanding end-to-end GeoAI, but that would be my direction for a business or research project.
See the GitHub repository for any updates or additional details on improvements.
Summary and future work
We have shown how TorchGeo is built on top of PyTorch and PyTorch Lightning to enable end-to-end GeoAI analyses.
We then used TorchGeo to run building detection using semantic segmentation with a U-Net neural network model.
While some GeoAI use cases can now be done in low- or no-code settings, it is helpful in many analyses to be able to return to the full code, while at the same time having available modularization and abstractions to make a project quicker to run and more organized than such code by itself.
The results here are reasonable, with most building detected and few false positives. Various possible detail improvements are listed in the text above.
Besides the improvements for this specific U-Net plus Inria buildings result, some other possible future work would be:
- Rerun in TorchGeo > 0.10.0 when out and its plots are fixed, instead of 0.7.1, giving more choices of metrics
- With better data and models, e.g., embeddings and modern geo foundation LLMs, plus building vector outputs, fine-tune a model to Millbrae buildings and improve upon existing Open Street Map buildings data
- Run other end-to-end workflows of interest in TorchGeo such as classification, object detection, or regression
