[论文解读] 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 Applications参考文献 11被引用 576

一句话总结

TernausNet 通过使用在 ImageNet 预训练的 VGG11 编码器来改进 U-Net，并比较三种权重初始化方案，显示更快的收敛和更高的验证 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.

研究动机与目标

• 激励并展示对编码器进行预训练如何使 U-Net 在像素级分割任务中的性能受益。
• 评估三种权重初始化方案并衡量它们对分割性能的影响。
• 展示在真实航空影像数据集上的收敛速度提升和最终 IoU。

提出的方法

• 用在 ImageNet 上预训练的 VGG11 编码器替换 U-Net 的编码器（不含全连接层）。
• 比较三种初始化方案：LeCun 均匀、在 ImageNet 预训练的 VGG11 编码器，以及在 Carvana 上完全预训练的模型。
• 在 Inria Aerial Image Labeling Dataset 上用 Adam 优化器训练 100 轮，并以 IoU 作为评估指标。
• 使用 0.3 的二值像素掩膜阈值从概率图中得到二值预测。

实验结果

研究问题

• RQ1 用在 ImageNet 上预训练的 VGG11 初始化 U-Net 的编码器，是否比随机初始化在分割性能上有提升？
• RQ2 在 Carvana 上的全网络预训练与在 ImageNet 预训练编码器的预训练相比，在 IoU 和收敛速度方面有何差异？
• RQ3 预训练是否能缩短训练时间并提升航空影像建筑物分割的最终验证 IoU？

主要发现

• 在 ImageNet 上进行的编码器预训练将 IoU 提升至 0.686，从随机初始化的 0.593 提升。
• 在 Carvana 上的全预训练获得 IoU 0.687，与 ImageNet 预训练的编码器相似。
• 预训练模型的收敛速度更快，达到的稳态 IoU 高于随机初始化模型。
• 该方法在 Inria Aerial Image Labeling Dataset 上进行演示，包含 150 张训练图像和 30 张验证图像。