Fine-Tune a Semantic Segmentation Model with a Custom Dataset

Summary

Fine-tuning a semantic segmentation model requires domain-appropriate training data and an efficient pipeline. Existing autonomous driving datasets feature roadway imagery captured by cars, creating a distribution mismatch for sidewalk-based delivery robots. To resolve this discrepancy, a dedicated dataset of sidewalk imagery is loaded from the Hugging Face Hub, split into training and test sets, and augmented on-the-fly using SegformerImageProcessor and torchvision. The smallest SegFormer architecture, B0, is fine-tuned using Hugging Face's Trainer API with mean Intersection over Union evaluation metrics. The final pipeline pushes the fine-tuned model to the Hub and executes inference by upsampling output logits to original image dimensions.

Takeaways

1. 1

   SegFormer employs a hierarchical Transformer encoder lacking positional encodings paired with a simple multi-layer perceptron decoder.

2. 2

   Applying dataset transforms on-the-fly via set_transform avoids the disk space overhead of preprocessing entire image datasets in advance.

3. 3

   SegFormer output logits have reduced spatial dimensions of height/4 and width/4, requiring bilinear interpolation to match original image dimensions before argmax classification.