Melanoma segmentation with U-Nets built from scratch in PyTorch

Problem

Outlining the border of a skin lesion precisely is the first step to analysing it: asymmetry and border irregularity are classic signs of melanoma. The goal is a model that, given a dermoscopic image, predicts the lesion mask pixel by pixel (lesion or background).

The dataset is ISIC Melanoma Segmentation: 2,750 images with their annotated masks. The project’s question is whether it pays off to start from an encoder pretrained on ImageNet, a very different domain, instead of training the whole network from scratch.

My contribution

I implemented a classic U-Net from scratch in PyTorch: five convolutional blocks with 64 to 1024 filters, a symmetric decoder with transposed convolutions and skip connections at every level (31 M parameters, all trainable).
I implemented a U-Net with a frozen pretrained ResNet-50 encoder, with skip connections from conv1, layer1, layer2 and layer3 and a custom five-block decoder. It trains 26.5 M of its 50 M parameters.
I wrote the full pipeline: loading and preprocessing, training with early stopping, threshold selection and evaluation.

Technical approach

Preprocessing512 × 512 px, greyscale, normalisation and binarised mask
TrainingAdam (1e-4) and BCE with logits, up to 50 epochsearly stopping: val. AUC
Optimal thresholdPrecision-recall curve on validation; the max-F1 threshold is chosen
Test evaluationPixel-level AUC-ROC and IoU (Jaccard index)

Pipeline shared by both models. Random 70/10/20 split: 1,924 training, 275 validation and 551 test images.

Images are converted to greyscale, so ResNet-50’s first layer, designed for RGB, does not fit. Instead of initialising it randomly, I averaged its weights across the three channels: it keeps what was learned about edges and textures and accepts a single channel.

# scripts/models/UNet_ResNet50_Encoder.py (excerpt)
resnet50 = models.resnet50(pretrained=True)
self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
# RGB to greyscale: average the weights of the three channels
self.conv1.weight.data = resnet50.conv1.weight.data.mean(dim=1, keepdim=True)

# Frozen encoder: only the decoder is trained
encoder = [self.conv1, self.bn1,
           self.layer1, self.layer2, self.layer3, self.layer4]
for layer in encoder:
    for param in layer.parameters():
        param.requires_grad = False

Two decisions in the evaluation protocol:

Sampled AUC. Computing AUC over every validation pixel does not fit in memory; it is computed on one pixel in ten, which is enough to drive early stopping.
Leakage-free threshold. The threshold that turns probabilities into a mask is chosen on validation (max F1 on the precision-recall curve) and then applied as-is on test. Choosing it on test would inflate the result.

Results

Model	Trainable parameters	AUC-ROC (test)	IoU (test)	Threshold
U-Net from scratch	31.0 M	0.9616	0.7236	0.42
U-Net with frozen ResNet-50	26.5 M	0.9769	0.7524	0.34

The pretrained encoder wins on both metrics (+1.5 AUC points and +2.9 IoU points) while training fewer parameters. ImageNet features, even from another domain, capture edges, textures and shapes that speed up convergence, and freezing them acts as regularisation with fewer than 2,000 training images.

Two test cases in three columns: original image, ground-truth mask and the mask predicted by the ResNet-50 U-Net. In the second case the prediction follows the irregular lesion border in more detail than the ground truth, which is a polygon. — ResNet-50 U-Net on test: image, ground-truth mask and prediction. The ground truth is a rough polygon; the prediction follows the irregular lesion border.

Two test cases for the U-Net trained from scratch: original image, ground-truth mask and predicted mask. The predictions cover large lesions well, with some holes and noise at the edges. — U-Net from scratch on test. It gets the overall shape of large lesions right, with more holes and edge noise.

Training curves of the ResNet-50 U-Net. Left: training and validation loss, falling from 0.33 to 0.14. Right: validation AUC, rising from 0.955 to a peak of 0.973 at epoch 26. — Training the ResNet-50 U-Net. Validation AUC peaks (0.973) at epoch 26 and early stopping ends training five epochs later.

Evidence

Repository with both architectures (scripts/models/) and the training notebook with all its outputs.

Professional application

Segmentation in medical imaging (lesions, organs, tumours) and other domains: defects in industrial inspection, plots in satellite imagery.
Transfer learning with little labelled data: reuse a pretrained encoder and train only what is specific to the problem.
Leakage-free evaluation protocols, with thresholds and model decisions always made outside the test set.