Source-linked AI summary
TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
Vladimir Iglovikov, Alexey Shvets
TL;DR
Accurate pixel-wise segmentation is important but difficult when labeled datasets are small. The paper equips U-Net with a VGG11 encoder initialized from ImageNet and compares initialization schemes, reporting faster convergence and higher steady-state values for pre-trained models. The architecture was also part of a first-place solution in the Carvana competition.
Problem
Pixel-wise segmentation requires reliable dense predictions for applications such as medical imaging and traffic safety, while segmentation datasets are often small and costly to label.
Method
The paper uses a U-Net encoder-decoder with a VGG11 encoder and compares LeCun initialization, ImageNet-pre-trained encoder weights, and Carvana-pre-trained full-network weights.
Results
Pre-trained models converged much faster and appeared to reach higher steady-state validation values than the non-pre-trained network.
Takeaways & Limitations
Pre-trained encoder weights improve U-Net performance and provide a practical initialization approach for segmentation with limited training data.
Abstract
from arXiv · showhide
Pixel-wise image segmentation is demanding task in computer vision. Classical U-Net architectures composed of encoders and decoders are very popular for segmentation of medical images, satellite images etc. Typically, neural network initialized with weights from a network pre-trained on a large data set like ImageNet shows better performance than those trained from scratch on a small dataset. In some practical applications, particularly in medicine and traffic safety, the accuracy of the models is of utmost importance. In this paper, we demonstrate how the U-Net type architecture can be improved by the use of the pre-trained encoder. Our code and corresponding pre-trained weights are publicly available at https://github.com/ternaus/TernausNet. We compare three weight initialization schemes: LeCun uniform, the encoder with weights from VGG11 and full network trained on the Carvana dataset. This network architecture was a part of the winning solution (1st out of 735) in the Kaggle: Carvana Image Masking Challenge.
I. INTRODUCTION
Image segmentation is a demanding pixel-level classification problem with important medical and safety applications. U-Net and pre-trained weights address the need for accurate segmentation with limited labeled data.
- Motivation: Reliable segmentation divides images into meaningful regions at the pixel level, supporting medical diagnosis and safe self-driving vehicles.Manual annotation is time-consuming and prone to mistakes and inconsistencies.
- Prior approaches: Fully Convolutional Networks replace fully connected layers with convolutions to produce spatial feature maps that are upsampled into dense outputs.This supports end-to-end segmentation for input images of arbitrary sizes.
- Prior approaches: U-Net extends this approach with skip connections that combine low- and high-level features for precise pixel-level localization.Its upsampling path also propagates context information through many feature channels.
- Contribution: The paper improves U-Net by using pre-trained weights, demonstrating the approach on aerial building segmentation and relating it to a winning Carvana solution.The Carvana solution ranked first among 735 teams.
II. NETWORK ARCHITECTURE
The model uses a U-Net encoder-decoder design with VGG11 as the encoder, progressive downsampling and upsampling, and skip connections for localization. ImageNet pre-training is motivated by the small size of typical segmentation datasets.
- U-Net structure: The contracting path captures context by downsampling feature maps while increasing channels, and the expansive path upsamples them through convolution.The two paths form the encoder-decoder structure of U-Net.
- Architecture overview: Figure 1 uses VGG11 without fully connected layers as the encoder, with feature-map height indicating spatial size and width indicating channel count.Channels increase through encoding stages and decrease through decoding stages.
- U-Net structure: Skip connections concatenate high-resolution contracting features with upsampled expansive features to improve localization, producing a pixel-by-pixel class mask.The expansive branch increases output resolution before combining features.
- Initialization: Typical segmentation datasets contain at most thousands of images, whereas training large networks without overfitting generally requires millions, motivating ImageNet initialization.Pre-training can reduce the number of layers that must be learned from the target dataset.
- VGG11 encoder: VGG11 contains 11 sequential layers, including seven convolutional layers with ReLU activations and five max-pooling operations.Each pooling operation halves the feature-map size.
- Encoder-decoder implementation: The encoder removes fully connected layers and uses a 512-channel convolutional bottleneck, while transposed convolutions double spatial size and halve channels in the decoder.Decoder outputs are concatenated with corresponding encoder features and convolved afterward.
- VGG11 encoder: Figure 2 identifies VGG11 feature-map channel counts in each box and places a ReLU activation after every convolutional layer.The figure provides the channel configuration used by the encoder.
III. RESULTS
The study evaluates U-Net segmentation on aerial building masks using Jaccard index and compares three weight-initialization schemes. Pre-trained models converge faster and appear to achieve higher steady-state performance than random initialization.
- Experimental setup: The Inria dataset contains 180 RGB aerial images labeled as building or not-building, with 30 images reserved for validation and 150 used for training.Images have 5000×5000 pixels, and training uses random 768×768 crops for 100 epochs.
- Evaluation and loss: Jaccard index measures similarity between predicted and reference pixel sets, while binary cross entropy is used as the binary-classification loss.The generalized loss combines the binary cross-entropy expression with the Jaccard formulation.
- Evaluation and loss: The output is converted to binary pixels using a validation-selected threshold of 0.3.Pixels below the threshold are set to zero, while pixels above it are retained.
- Experimental setup: The experiment trains three otherwise identical U-Nets for 100 epochs, differing in LeCun uniform, ImageNet-pre-trained VGG11 encoder, or Carvana-pre-trained full-network initialization.The decoder remains LeCun-initialized for the VGG11 condition.
- Results: Pre-trained models converge much faster to a steady value than the non-pre-trained network, and their steady-state value appears higher.Figure 3 compares random initialization, an ImageNet-pre-trained VGG11 encoder, and full-network Carvana pre-training.
- Results: After 100 epochs, the three initialization schemes produce visibly different prediction quality in the aerial building masks.Figure 4 compares the ground-truth mask with predictions from random, VGG11-encoder, and Carvana-pre-trained models.
IV. CONCLUSION
The paper concludes that fine-tuning an encoder with pre-trained weights improves U-Net performance for image segmentation. It also identifies stronger encoders and open-source implementation as directions and resources for continued use.
- Conclusion: Fine-tuning pre-trained encoder weights improves U-Net performance and is especially suitable when segmentation datasets are costly to collect and label.The authors also state that pre-trained networks substantially reduce training time and help prevent over-fitting.
- Conclusion: The approach can be extended with pre-trained VGG16 or ResNet-family encoders while retaining a simple decoder.The paper’s code is released under an MIT license.