Abstract

We describe the diagnosis of a subtle off-by-one bug in an image segmentation pipeline built on segmentation-models.pytorch (SMP), used to detect Retrogressive Thaw-Slumps in satellite imagery. Large satellite scenes are cropped into overlapping tiles and stitched back together with a triangular-weighted blending scheme, but predictions from certain encoder-decoder architectures were occasionally shifted by one pixel toward the center of their tile. Through toy-dataset reproductions and controlled retraining across four architecture-encoder combinations, we trace the shift to an internal size mismatch in SMP's decoder: rather than rescaling feature maps to the expected output size, the library silently pads them, introducing a directional one-pixel offset near tile borders. We document the debugging process, the workaround we applied, and the broader lesson that even mature, widely used computer vision libraries can hide small but consequential bugs.

Introduction

Machine Learning algorithms can be truly magical. Thanks to today's universe of righ Python libraries, with a little bit of data and a correctly applied method one can create very well working models without problems. Some of the methods are bigger and more complex than others, especially in computer vision, sometimes leaving one with an almost-black-box, making debugging especially hard. This is the story of how I spent a complete week huntig a suspicouos and hidden bug in one of the models of my most beloved image segmentation framework: segmentation-models.pytorch(SMP).

My team and I try to find so called Retrogressive Thaw-Slumps (RTS) on satellite images, which is a distinct landscape feature of permafrost regions. For this we created an image segmentation pipeline, utilzing multiple optical bands from the satellite imagery and some additional auxiliary datasets, e.g. elevation data. The core of our segmentation library are the highly configurable models of the SMP library. We trained multiple different encoder-decoder architectures, encoder implementations until we found reasonably well working combinations of them. As any respected and modern computer vision ML scientist, we started with the gloryfied U-Net in combination with a Res-Net encoder. This worked well, later we also tried out the transformer based MaxViT encoder with the novel UPerNet architecture, which worked slightly better on our metrics and also resulted in better results visually. From there on we used this transformer based combination our further research.

A big problem when working with geo-spatial data and computer vision models is the gigantic size of satellite imagery. They often come in 10.000 x 10.000 pixel, which is way to large for any SMP model to handle, so it needs to be cropped into smaller parts. To reduce edge-effects at the places where the larger imagery is cropped, we use a triangular overlap technique, where the smaller tiles are cropped-out with a specified overlap so the area around a crop is predicted twice. Then the probabilities are merged together, weighted by how far a pixel is away from the border, leaving with normalized probability layer with no cuts visible. At least in theory.

When one member of my team did some deep-dive analysis on the predictions from the new UPerNet-MaxViT model, they found out that the outputs are sometime shifted by one pixel in a random direction. This didn't occured on every prediction, only some where affected. When he first showed me this, I thought it is a model intrinsic in-perfection, some un-debuggable mis-trained inner state. However, he convinced me to have a closer look on it, since he didn't trusted my triangular overlap implementation. Neither did I, so I started diving into it. After two days, several new unit tests and a partial rewrite of my implementation, I was very sure my implementations, the first and the new one, are correct. I created multiple visualizations on a toy dataset to back up my bold claims. And they turned out to be correct.

Toy Pseudo Code

img = random_noise()
model = (x) => 2 * x  # deterministic
for patch in patches(img):
    out[patch-idx] = model(patch)
final = merge(out)
assert final == img * 2

At day three I had luck and found an important relationship: The shifts only occured on the edge of a tile and always into the direction pointing to it's center. This was a turning point, as this indicated that the bug may be model related due to the several rescalings done by them. My thesis was that the novel UPerNet outputs were two pixels too small, so it needed to rescale it, but instead of rescaling they instead pad it. So I trained a models on all four arch-encoder combinations from scratch on our real training data, to see if it is really the architecture and not an encoder or general problem. And it turned out my thesis was correct. Then I started digging into the code of SMP and found this:

    ...
    # UPerNets Segmentation head
    self.segmentation_head = SegmentationHead(
        in_channels=decoder_channels,
        out_channels=classes,
        activation=activation,
        kernel_size=1,
        upsampling=upsampling,  # defaults to 4
    )
    ...

    class SegmentationHead(nn.Sequential):
        def __init__(
            self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
        ):
            conv2d = nn.Conv2d(
                in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
            )
            upsampling = (
                nn.Upsample(mode="bilinear", scale_factor=upsampling, align_corners=True)
                if upsampling > 1
                else nn.Identity()
            )
            activation = Activation(activation)
            super().__init__(conv2d, upsampling, activation)

The UPerNet implementation of SMP indeed applies an upsampling per default. To workaround the bug, we manually rescaled and cropped the outputs to the right size before further merging them together, which isn't a perfect solution but it works for now. As I have not yet run a deeper experiments with it, I have not opened an issue on the SMP repository yet. Indicated by the upsampling default of 4 (which is probably a very sane default), our problem was probably at OSI layer 8: a wrong model argument.

I find it incredible how nuanced and hidden some bugs can be. Every now and then I encounter such niche problems and they always leave me wondering: how many of these bugs are out there, hidden in the shadows, waiting for us to run into them.