Fine-Tuning SegFormer for Semantic Segmentation with Custom Datasets

Hugging Face has detailed a workflow for fine-tuning SegFormer, a state-of-the-art semantic segmentation model, to handle custom image datasets. This process enables the creation of specialized models—such as one designed for a pizza delivery robot to recognize sidewalks and obstacles—by leveraging pre-trained weights and the Hugging Face ecosystem.

Understanding SegFormer and Semantic Segmentation

Semantic segmentation is the process of classifying every individual pixel in an image, providing a more granular level of detail than standard image classification. This capability is critical for applications like medical imaging and autonomous driving, where precise boundary detection (e.g., identifying the exact edge of a sidewalk) is required.

SegFormer, introduced by Xie et al. in 2021, improves upon previous convolutional neural network (CNN) and Vision Transformer (ViT) approaches. Its architecture consists of:

  • A Hierarchical Transformer Encoder: Unlike ViT, SegFormer does not use positional encodings.
  • A Simple MLP Decoder: A multi-layer perceptron decoder that processes the encoder's output.

Data Preparation and Augmentation

Effective semantic segmentation requires datasets with precise segmentation maps. While general datasets like ADE20k, CityScapes, or BDD100K exist, domain-specific data is often necessary to avoid distribution mismatch. For instance, a robot operating on sidewalks requires data captured from a sidewalk perspective rather than a car's perspective.

Dataset Loading and Processing

Using the datasets library, custom datasets (such as segments/sidewalk-semantic) can be loaded and split into training and testing sets. To ensure the model receives data in the correct format, the SegFormerImageProcessor is used.

On-the-Fly Transforms

To optimize disk space and training speed, Hugging Face recommends using transforms via set_transform. This prepares batches of data on-the-fly rather than preprocessing the entire dataset in advance. To increase model resilience to varying lighting conditions, torchvision.transforms.ColorJitter is integrated into the training pipeline to randomly adjust brightness, contrast, saturation, and hue.

Fine-Tuning Workflow

Model Selection

SegFormer offers five model sizes (B0 through B5). For edge deployment—such as on a delivery robot—the B0 model is recommended due to its small footprint (approximately 14MB) and efficiency. The fine-tuning process typically starts with a model pre-trained on ImageNet-1k (nvidia/mit-b0).

Training Configuration

Fine-tuning is managed via the Hugging Face Trainer API. Key configuration parameters include:

  • Hyperparameters: Learning rate (e.g., 0.00006) and epoch count (e.g., 50).
  • Evaluation Metric: The mean Intersection over Union (mIoU) is used to measure the overlap between the predicted segmentation mask and the ground truth.
  • Logit Upscaling: Because SegFormer outputs logits at 1/4 the original image resolution (height/4, width/4), they must be upscaled using bilinear interpolation to match the label size before computing mIoU.

Inference and Deployment

Once fine-tuned, the model and its image processor can be pushed to the Hugging Face Hub. This allows for easy sharing and the creation of an inference widget for real-time testing via the hosted inference API.

Executing Inference

To perform inference on a new image, the following steps are required:

  1. Preprocessing: Process the image using SegformerImageProcessor.
  2. Forward Pass: Pass the processed image through the model to obtain logits.
  3. Rescaling: Upsample the logits to the original image dimensions using nn.functional.interpolate.
  4. Prediction: Apply an argmax operation on the class dimension to determine the final pixel-level category predictions.

Sources