[Paper Review] TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation

Vladimir Iglovikov, Alexey A. Shvets|arXiv (Cornell University)|Jan 17, 2018

Advanced Neural Network ApplicationsComputer Science11 references576 citations

TL;DR

TernausNet improves U-Net by using a VGG11 encoder pre-trained on ImageNet and compares three weight initialization schemes, showing faster convergence and higher validation IoU.

ABSTRACT

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.

Motivation & Objective

• Motivate and demonstrate how pre-training an encoder benefits U-Net for pixel-wise segmentation tasks.
• Evaluate three weight initialization schemes and measure their impact on segmentation performance.
• Showcase improved convergence speed and final IoU on a real aerial image dataset.

Proposed method

• Replace the U-Net encoder with a VGG11 encoder (without fully connected layers) pre-trained on ImageNet.
• Compare three initialization schemes: LeCun uniform, VGG11 ImageNet-pretrained encoder, and fully pre-trained Carvana model.
• Train with 100 epochs on the Inria Aerial Image Labeling Dataset using Adam optimizer and measure IoU as the evaluation metric.
• Use a binary pixel mask threshold of 0.3 to obtain binary predictions from probability maps.

Experimental results

Research questions

• RQ1 Does initializing the U-Net encoder with a pre-trained VGG11 on ImageNet improve segmentation performance compared to random initialization?
• RQ2 How does full network pre-training on Carvana compare to ImageNet-pretrained encoder pre-training in terms of IoU and convergence speed?
• RQ3 Can pre-training reduce training time and improve final validation IoU for aerial image building segmentation?

Key findings

• Encoder pre-training on ImageNet improves IoU to 0.686 from 0.593 with random initialization.
• Full pre-training on Carvana yields IoU of 0.687, similar to ImageNet-pretrained encoder.
• Pre-trained models converge faster and reach higher steady-state IoU than randomly initialized models.
• The approach is demonstrated on the Inria Aerial Image Labeling Dataset with 150 training images and 30 validation images.