Skip to main content

A U-Net ablation study for retinal fluid segmentation in OCT

00:03:04:26

Optical Coherence Tomography (OCT) produces cross-sectional scans of the retina, and one of the things clinicians look for in them is fluid that has leaked into layers where it shouldn't be — a sign of retinal disease that needs monitoring over time. Segmenting that fluid automatically, pixel by pixel, is a fairly standard semantic segmentation problem, but with a small, expensive-to-label medical dataset it's also a good candidate for asking a more useful question than "what's the best model": which specific design choices actually move the needle when you only have 50 image-mask pairs to work with?

The setup

The dataset was 50 OCT image/mask pairs — small enough that a single train/test split would have made the results mostly noise. I used 5-fold cross-validation instead, training and evaluating five times over different splits and reporting the average, which gives a much more honest picture of how a configuration generalises rather than how lucky one split happened to be.

All six configurations shared the same U-Net encoder/decoder skeleton (built with segmentation-models-pytorch) and varied along two axes:

  1. Encoder — a plain encoder trained from scratch, versus a ResNet34 backbone pretrained on ImageNet.
  2. Loss function — Binary Cross-Entropy (BCE) alone, Dice loss alone, or a combined BCE + Dice loss.
python
# Combined BCE + Dice loss used for the best-performing configuration
import torch.nn as nn
import segmentation_models_pytorch as smp

bce = nn.BCEWithLogitsLoss()
dice = smp.losses.DiceLoss(mode='binary')

def combined_loss(pred, target):
    return bce(pred, target) + dice(pred, target)

What actually moved the needle

The best configuration — ResNet34 pretrained encoder with the combined BCE + Dice loss — reached a Dice coefficient of 0.824 ± 0.010, an IoU of 0.716, and an AUC-ROC of 0.993 across the 5 folds.

A few things came out of comparing all six configurations against each other:

  • Pretraining beat training from scratch by a wide margin, which isn't surprising given only 50 samples — a from-scratch encoder simply doesn't see enough data to learn useful low-level features, whereas the ImageNet-pretrained weights already encode general edge and texture detectors that transfer reasonably well to grayscale medical imagery.
  • The combined loss consistently beat either loss alone. BCE alone is a per-pixel loss that doesn't care about the shape or connectivity of the segmented region, so it can look fine on the loss curve while still producing fragmented masks. Dice alone directly optimises for mask overlap but its gradient can be unstable at the start of training when predictions are close to empty. Adding them together tempered that instability while still pushing the model toward better overlap.
  • Augmentation mattered more for the from-scratch encoder than the pretrained one. With so little data, elastic transforms and colour jitter gave the from-scratch model more effective variety to learn from, but the pretrained encoder was already regularised enough by its pretraining that augmentation helped less.

Why cross-validation, not a single split

With only 50 samples, a single 80/20 split leaves the test set at 10 images — small enough that one or two unusually hard (or easy) cases can swing the reported Dice score by several points and make one configuration look better than another purely by chance. Running all six configurations through the same 5-fold split let me compare average performance and its variance, which is what actually separated a genuinely better configuration from a lucky split.

Takeaway

None of the individual pieces here are exotic — transfer learning and combined loss functions are well-known techniques. What the ablation actually provided was evidence for which combination of well-known techniques mattered for this dataset, instead of assuming it based on general intuition. That's usually the more useful output of a small-data medical imaging project: not a novel architecture, but a validated recipe.

Code and full results are on GitHub, and the project page has more detail on the architecture and evaluation setup.