TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation

guthub：TernausNet
Unet

数据集：Kaggle Carvana

摘要

像素级的语义分割，经典网络 U-Net，由编码encoders和解码decoders组成，是医学图像，卫星图像的分割常用的网络结构. 通常情况下，神经网络的权重初始化是采用在 ImageNet 上预训练的权重参数进行的，相比于重新训练，在小数据集上的效果更佳. TernausNet 阐述了如何采用预训练的编码器encoders 来改进 U-Net，以提升图像分割精度. 对比了三种权重初始化方案：

• LeCun uniform
• VGG11 的编码器
• 完全在 Carvana 数据集训练的网络

介绍

CNN在视觉领域有很好的表现。

FCN：目前很好。这种方法的主要思想是使用CNN作为一个强大的特征提取器，通过卷积代替完全连接层来输出空间特征图，而不是分类的分数。这些特征图被进一步上采样，以产生密集的像素输出。这种方法允许以端到端的方式训练CNN，以便用任意大小的输入图像进行分割。准确率相比于之前比较常用的方法，有提高

U-Net architecture

U-Net 是一种编码器-解码器结构，在网络向前传播的过程中，编码器逐渐收缩，减少池化层的空间维度；而解码器逐渐扩张，逐步修复物体的细节和空间维度。编码器和解码器之间通常存在跳跃连接（skip connection），跳跃连接能够将低层次的特征图和高层次的特征图结合起来，能帮助解码器更好地修复目标的细节，实现像素级的定位。在上采样部分，大量的特征通道能向更高分辨率的层传送上下文信息。
适用于二值分割

网络结构

改进后的框架：TernausNet

通常情况，U-Net 采用随机初始化权重重新开始训练。

改进点：在 ImageNet 上预训练的来初始化网络权重。通过这种方式，学习过程可以针对网络的非预先训练的几个层(有时仅针对最后一层)进行，以考虑数据集的特征。

骨架（编码器）：经过改进的VGG11

此网络编码器部分采用VGG11（VGG11包含7个3x3卷积层，每个层后加一个ReLU，同时进行5次最大池化操作，具体如上图）

改进点：为了达到语义分割的目的，编码器采用了类似 FCN 的设计思路，将 VGG11 的最后 3 层替换为了一个 512 通道的卷积层，它同时也是编码器与解码器的「瓶颈」，将 U-Net 的左右两个部分分开。

解码器

这里采用转置卷积层放大特征图尺寸并且减少一半原通道数。同时将转置卷积的输出与解码器的相应部分的输出串联（跳跃连接）。特征图的结果通过卷积操作使得其通道数与对应编码器部分相同。这一上采样过程重复5次对应5次池化操作。

传统全连接层可接受任意大小图片输入，但因为此处有5个池化层，每次图像缩小到原来一半，即缩小2^5=32倍，因此当前网络要求输入图像大小需要能被32整除

实验结果

优化损失函数

Inria Aerial Image Labeling Dataset 数据集包含 180 张欧洲和美国的城市住区的航空图像，并被标注为建筑(building) 和非建筑(not building) 两类. 每张图片是 5000x5000 分辨率的 RGB 图像，每个像素对应于地球表面的 30x30 平方厘米. 30 张作为验证集，150 张作为训练集，训练了 100 个 epochs. 随机裁剪为 768x768 作为训练输入，中心裁剪 1440x1440 用于验证输入。

采用 Jaccard index(IoU，Intersection Over Union) （相似系数）作为评估度量. 两个数据集 A 和 B 的相似性度量的 IoU 定义如下：

在图像分割任务中，由于图像由像素点组成，在离散问题中，我们可以将jaccard相似系数改写为：
其中，上面没有三角形的是gt，有三角形的是通过模型预测得出的像素点的概率。

此时，可以将图像分割问题看作一个像素的二分类问题，其交叉熵损失函数可以表示为：

整个 TernausNet 的损失函数可以表示为：

最小化这个损失函数，我们可以同时最大化要预测的右边的像素的概率，并且最大化掩模（真实值）和相应预测之间的交集。

后处理

为了得到二值像素值，采用 threshold=0.3 进行二值化. 根据不同数据集进行选择该阈值. 所有小于阈值的像素设为 0，大于阈值的像素设为 1. 然后每个像素值乘以 255，即可得到最终的黑白预测mask. 不同权重初始化网络的结果：