原文链接:https://www.yuque.com/yahei/hey-yahei/segmentation
简单过一下语义分割的主流框架——FCN、UNet、SegNet、PSPNet、DeepLab
分割任务论文集与各方实现:https://github.com/mrgloom/awesome-semantic-segmentation
pytorch model zoo:https://github.com/yassouali/pytorch_segmentation
gluon model zoo:https://gluon-cv.mxnet.io/model_zoo/segmentation.html
SOTA Leaderboard:https://www.paperswithcode.com/task/semantic-segmentation
论文:《Fully Convolutional Networks for Semantic Segmentation(CVPR2015)》
参考:《FCN的学习及理解 | CSDN, moonuke》
主要贡献:
为了跟图示统一,后续将以AlexNet 为backbone进行讨论。
就像《卷积神经网络CNN - 全连接层 | Hey~YaHei!》所提到的,卷积神经网络CNN因为全连接层的限制,要求网络输入具有固定的尺寸大小。FCN作者将最后的三个全连接层换成1x1卷积,如果输入特征图恰好是1x1,那明显是等价的;如果不是1x1,那网络也不至于出错,此时输出大小由输入大小决定。
FCN用反卷积(deconvolution)在网络深层做上采样操作,以恢复出输入图片同等尺寸的分割结果,也就是每个像素点的类别。
反卷积其实相当于零填充上采样+卷积,与padding不同的是,它填充在输入特征图的像素点之间。
3x3标准卷积 | 3x3反卷积 | 3x3空洞卷积 |
---|---|---|
| (以上三图,蓝色方块为输入特征图的像素点,绿色方块为输出特征图像素点,空白部分填零)
| | |
关于反卷积的详细过程此处不再赘述,感兴趣可以参考《怎样通俗易懂地解释反卷积? | 知乎, 孙小也》
值得一提的是,通常部署的时候不喜欢用反卷积,因为推理引擎往往没有针对反卷积做充分的优化。大多都直接用双线性插值/三线性插值做上采样(简化操作顺便提高推理速度),顶多再叠一层普通卷积来进一步提取特征。
众所周知,浅层特征注重细致的局部、位置信息,深层特征注重抽象的全局、分类信息。分类任务里不关注位置信息,所以随着网络前传,即使特征图分辨率越来越小,信息越来越抽象,位置信息逐步丢失,也无伤大雅。但检测任务和分割任务不同,除了需要给出对象的分类之外,还得给出位置信息——因此浅层特征的位置得想办法把它保留下来,比较直观的想法就是跨层把浅层和深层信息融合起来。
融合的方式有很多,最简单的如逐元素相加/相乘,或者连接特征(一般是从通道维度上做拼接)后做进一步的特征提取。FCN采用的即是简单的逐元素相加的形式,以500x500x3的输入图片为例(虚线以上就一个普通的全卷积网络,虚线以下是跨层融合相关的层)——
FCN采用普通的softmax交叉熵作为损失函数,既然通道方向决定了每个像素点的类别,那就对每个像素点计算softmax交叉熵,最后加和起来作为最终的损失进行训练。
原文倒是讲究,采用分阶段训练的方式,用与训练好的分类网络作为backbone,丢弃最后一层全连接,其他全连接替换成卷积并重新初始化权重(丢弃原有的全连接权重),再逐一融合中间层特征进行多阶段训练。
阶段 | 训练部分 |
---|---|
#1 | |
#2 | |
#3 | |
#4 |
论文:《U-Net: Convolutional Networks for Biomedical Image Segmentation(MICCAI2015)》
主要贡献:
网络整体结构呈现U型的对称结构,故称为UNet,左半部分为卷积组成的下采样路径,右半部分为反卷积和卷积组成的上采样路径。下采样路径每个阶段包含两_次_卷积和一次池化(卷积通道扩增+卷积等通道特征提取+池化下采样),上采样路径每个阶段包含一次反卷积和两次卷积(反卷积上采样并收缩通道+卷积通道收缩+卷积等通道特征提取)。
与FCN不同的是,UNet采取的是拼接(Concat)+进一步提取特征的特征融合方式,如上图上采样路径最左侧的蓝色外框空白填充的方块。相比粗暴的逐点相加,拼接更好的保留浅层的特征信息,但相对地也会增加计算的开销。
为了保证拼接的正确性,浅层特征图需要裁剪到目标尺寸,如上图下采样路径最右侧蓝色方块的虚线部分,这里通常采用的是中央裁剪。拼接后由于通道倍增,按照**“若特征图尺寸倍增,则通道减半;若尺寸减半,则通道倍增”**的惯例,需要先用卷积把通道数收缩到原来的一半。
UNet对大尺寸图像分割任务采用了重叠切片的平铺策略。
首先要注意到UNet与常规的网络不同,所有的卷积和池化都不加以padding(为了2x2池化下采样不加padding由不丢失信息,需要保证每次池化输入的尺寸为2的倍数),于是每做一次卷积,特征图都会稍微收缩一点点,这就导致了UNet的输出尺寸小于输入尺寸(如上图,572x572的输入最后出来只有392x392的掩膜)。
不padding意味着不引入无效信息,直观上是有好处的;另外这相当于用一张更大的图像来预测中央小区域的分割结果,相当于在分割的时候输入了目标区域以外的外围信息(如上图,蓝色框表示输入UNet的图像,最终只能产生黄色框部分的分割结果,实际在推理黄色框分割结果的时候也引用了黄色框以外蓝色框以内的外围信息的),这有助于提高模型的表现。
此外,UNet每步迭代采用单张图片输入,通过最大化输入图片尺寸来充分利用显存,同时对优化器采取一个较大的动量(如0.99)使之前的迭代结果能对本次迭代产生较大的影响,以此稳定训练过程。
注意:为了跟其他框架统一,yassouali/pytorch_segmentation依旧做了padding来保证输出图跟输入图尺寸一致。
首先看一下softmax交叉熵:
E ( x ) = ∑ l o g ( p ( x ) ) = ∑ l o g ( e x i ∑ j = 1 N e x j ) E(x) = \sum log(p(x)) = \sum log(\frac{e^{x_i}}{\sum^N_{j=1} e^{x_j}} ) E(x)=∑log(p(x))=∑log(∑j=1Nexjexi)
为了着重某些特殊的像素点,可以赋予一个权重 ω ( x ) \omega(x) ω(x),此时损失函数改造为
E ( x ) = ∑ ω ( x ) l o g ( p ( x ) ) E(x) = \sum \omega(x) log(p(x)) E(x)=∑ω(x)log(p(x))
而
ω ( x ) = ω c ( x ) + ω 0 ⋅ e x p ( − ( d 1 ( x ) + d 2 ( x ) ) 2 2 σ 2 ) \omega(x) = \omega_c(x) + \omega_0 \cdot exp(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2}) ω(x)=ωc(x)+ω0⋅exp(−2σ2(d1(x)+d2(x))2)
其中,
ω c ( x ) \omega_c(x) ωc(x)是权衡分类为每个分类所设置的一个损失权重;
ω 0 \omega_0 ω0和 σ \sigma σ也是人为设置的权重,如论文中推荐的 ω 0 = 10 , σ ≈ 5 \omega_0=10, \sigma \approx 5 ω0=10,σ≈5
d 1 ( x ) d_1(x) d1(x)和 d 2 ( x ) d_2(x) d2(x)分别代表当前像素点到最近和次近的细胞的欧式距离
将计算出来的 ω ( x ) \omega(x) ω(x)可视化后可以得到上图中的(d)
由于损失中对边缘位置的像素作出的较重的惩罚,最终将鼓励网络在实例边界做出较好的区分。
原文中还提到了对卷积参数采用标准差为 2 / N \sqrt{2/N} 2/N的高斯分布初始化方式(如3x3卷积,输入通道数为64,有 N = 3 × 3 × 64 N=3 \times 3 \times 64 N=3×3×64),这种方式其实跟He Initialization 也差不多。事实上自从BN层 的出现之后,深度学习网络对参数初始化也不再那么敏感。
UNet变种:《图像分割的U-Net系列方法 | 知乎, taigw》
UNet++:《研习U-Net | 知乎, 周纵苇》
论文:《SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation(TPAMI2016)》
参考:《SegNet图像分割网络直观详解 | 知乎, 郭冠华》
主要贡献:用反池化替代反卷积进行上采样,简化上采样过程,降低计算开销
整体结构跟FCN和UNet其实差不多,主要差别在于上采样的手段变成了反池化。
每个滑窗只会采样最大值作为输出,反池化上采样则是反过来,把一个像素值填到一个2x2的输出框内,为了跟下采样对应,需要在做最大值池化的时候记录采样点的索引(如总体框架图上的Pooling Indices信息),反池化的时候则填到对应位置上的。其余三个像素点则直接填零,由后续的卷积层完成特征图的平滑处理。
反池化有三个优点:
pytorch中的MaxPool可以直接返回索引(加入参数 return_indices=True
即可)并且直接提供了反池化的接口,所以反池化的实现也不难,具体参考pool的定义(L25)和使用(L69)、unpool的定义(L48)和使用(L88)。
论文:《Pyramid Scene Parsing Network (CVPR2017)》
参考:《论文笔记:Pyramid Scene Parsing Network | 简书, Efackw13》
《【图像分割模型】多感受野的金字塔结构—PSPNet | 知乎, 言有三》
主要贡献:
作者在场景解析的实际任务当中发现FCN存在以下问题:
为了解决上述三个问题,PSPNet提出了金字塔池化模块,该模块插入在输出分割结果前的最后一个特征图后边。
首先由CNN提取出特征图(PSPNet没像FCN, UNet, SegNet一样建立浅层到深层之间的跨层连接),然后经过不同的池化层下采样出不同尺寸的特征图(原文下采样为1x1, 2x2, 3x3, 6x6四种特征图),接着分别由卷积层将通道收缩为原来的1/N(原文中N=4)以保证拼接之后通道数与原来相同,由此得到不同尺寸的感知区域的局部信息。再将不同尺寸的局部信息上采样为原来特征图的尺寸但不改变通道数量(原文用双线性插值来上采样),与原始特征图拼接起来,最后经过卷积层映射到目标空间得到分割结果。
参考pytorch的定义yassouali/pytorch_segmentation/models/pspnet.py#L11-L38
以hszhao/PSPNet | github 为例,用netron可视化后可以看到详细的金字塔池化模块及后续处理的结构:
除了最终的分割分类损失之外,PSPNet还在中间位置加入了辅助损失,如下图所示,对ResNet第四阶段的输出特征图提前取出并且上采样到输入图片的尺寸,然后计算辅助损失loss2,并与主损失loss1加权求和后反传。
如yassouali/pytorch_segmentation/trainer.py#L61 | github 设置了权重为0.4;
用于计算loss2的特征图的产生,具体也可以参见 yassouali/pytorch_segmentation/models/pspnet.py 的L65-L71和L90-L94
参考:《【语义分割系列:一】DeepLab v1 / v2 论文阅读翻译笔记 | CSDN, 鹿鹿最可爱》
《【语义分割系列:五】DeepLab v3 / v3+ 论文阅读翻译笔记 | CSDN, 鹿鹿最可爱》
《deeplab系列总结(deeplab v1& v2 & v3 & v3+) | CSDN, Dlyldxwl》
论文:《Semantic Image Segmentation with Deep Convolutional Nets and Fully Connected CRFs (ICLR2015)》
主要贡献:
(为了与原文对应,以下讨论均以VGG16为backbone)
_
(缓解第一个问题)
按照常规利用分类网络做backbone的方式,作者剥离掉VGG16最后的三层全连接,此时最后一层卷积层的输出特征图分辨率为7x7,与输入的原始图像分辨率224x224相比,已经下采样了32倍,丢失了非常多的细节信息(原文中称之为很sparse、不dense)。
为了保留更多的信息,作者将最后两层池化的步长修改为1,也即取消了这两个池化的下采样功能,此时相当于只下采样了8倍。但如上一小节所说,这样就带来新的感受野偏小的问题。
DeepLab使用空洞卷积(dilated convolution)来解决这个问题,同时节约了不少的计算量。关于空洞卷积的图示和几种卷积的比较可以参考 FCN的反卷积上采样 一节。
(“下采样-标准卷积-上采样”和“空洞卷积”的效果比较,图源自v2)
(“标准卷积-下采样”和“空洞卷积”特征图尺寸变化示意,图源自以ResNet为backbone的v3)
简单来说,**空洞卷积通过跳跃性地采样,能以3x3的卷积核得到5x5、7x7等更大的等效感受野。**前述两个取消下采样功能的池化层之后的卷积层就换成了成空洞卷积。
注意到空洞卷积对输入点的采样是违背数据局部性原则的,如果没做些特殊处理那么执行的效率会很低(比如据说pytorch的空洞卷积速度就只有标准卷积的十分之一)。为了稍微高效点地实现空洞卷积,通常在底层实现上会将输入特征图按间隔采样分成若干输入图,依次执行标准卷积后再合并成最终空洞卷积的结果。
但无论如何,空洞卷积的执行效率都会明显低于标准卷积,尤其是在一些不作优化(空洞卷积本身也难优化)的框架上,所以在应用层面上来看,其实并不建议使用空洞卷积来设计网络。
缓解了第一个问题后,论文同时指出随后的上采样不再需要反卷积来恢复分辨率,直接双线性插值就可以得到可观的结果。训练时直接对ground truth下采样8倍,然后与改造后的VGG16输出求交叉熵作为损失函数;预测时则直接双线性插值得到分割的结果。
(缓解第二个问题)
全连接条件随机场(Fully Connected Conditional Random Field, Fully Connected CRF)
如上图的 DCNN output 一列所示,经过层层的下采样和上采样,特征图逐渐丢失部分信息,导致最终输出的图像显得比较平滑。而分割任务希望分割的结果边缘轮廓能够比较清楚犀利,于是作者引入了全连接CRF对CNN的输出结果进行后处理。
在传统的图像处理中,CRF通常用相邻像素来设计能量函数,从而消除一些噪音,达到平滑处理的目的。然而在分割中,我们的目标是恢复局部信息而非进一步平滑处理。
因此作者借鉴了《Efficient Inference in Fully Connected CRFs with Gaussian Edge Potentials (NIPS2011)》全连接条件随机场来实现分割结果的锐化处理。其能量函数为
E ( x ) = ∑ i θ i ( x i ) + ∑ i j θ i j ( x i , x j ) E(\boldsymbol{x})=\sum_{i} \theta_{i}\left(x_{i}\right)+\sum_{i j} \theta_{i j}\left(x_{i}, x_{j}\right) E(x)=i∑θi(xi)+ij∑θij(xi,xj)
其中
{ θ i ( x i ) = − log P ( x i ) θ i j ( x i , x j ) = μ ( x i , x j ) ∑ m = 1 K w m ⋅ k m ( f i , f j ) \begin{cases} \theta_{i}\left(x_{i}\right)=-\log P\left(x_{i}\right) \\ \theta_{i j}\left(x_{i}, x_{j}\right)=\mu\left(x_{i}, x_{j}\right) \sum_{m=1}^{K} w_{m} \cdot k^{m}\left(\boldsymbol{f}_{i}, \boldsymbol{f}_{j}\right) \end{cases} { θi(xi)=−logP(xi)θij(xi,xj)=μ(xi,xj)∑m=1Kwm⋅km(fi,fj)
对于一元项 θ i ( x i ) \theta_i(x_i) θi(xi)来说, P ( x i ) P(x_i) P(xi)是CNN产生的像素点 i i i的概率分布;
对于二元项 θ i j ( x i , x j ) \theta_{ij}(x_i, x_j) θij(xi,xj)来说,
μ ( x i , x j ) = { 1 , x i ≠ x j 0 , x i = x j \mu(x_i, x_j) = \begin{cases} 1, & x_i \neq x_j \\ 0, & x_i = x_j \end{cases} μ(xi,xj)={ 1,0,xi=xjxi=xj,这意味着每个像素点都会和全图像的所有像素点建立联系,也即“全连接”;
求和项里是 w m w_m wm加权的应用在像素特征 f i , f j \boldsymbol{f}_{i},\boldsymbol{f}_{j} fi,fj上的高斯核函数,原文采用像素点值和位置构造核函数
w 1 exp ( − ∥ p i − p j ∥ 2 2 σ α 2 − ∥ I i − I j ∥ 2 2 σ β 2 ) + w 2 exp ( − ∥ p i − p j ∥ 2 2 σ γ 2 ) w_{1} \exp \left(-\frac{\left\|p_{i}-p_{j}\right\|^{2}}{2 \sigma_{\alpha}^{2}}-\frac{\left\|I_{i}-I_{j}\right\|^{2}}{2 \sigma_{\beta}^{2}}\right)+w_{2} \exp \left(-\frac{\left\|p_{i}-p_{j}\right\|^{2}}{2 \sigma_{\gamma}^{2}}\right) w1exp(−2σα2∥pi−pj∥2−2σβ2∥Ii−Ij∥2)+w2exp(−2σγ2∥pi−pj∥2)
这里包含两项,第一项包含位置信息 p p p和值信息 I I I,后一项只考虑位置信息 p p p,两者通过 w 1 w_1 w1和 w 2 w_2 w2加权;
而 σ α , σ β , σ γ \sigma_\alpha, \sigma_\beta, \sigma_\gamma σα,σβ,σγ也是人工设置的超参数
和FCN、UNet一样,作者也尝试从浅层抽取特征与深层融合完成最后的分割预测。
具体来说,在输入图片和中间池化结果(前四个池化层的输出)上分别加两层卷积(3x3+1x1)做特征提取和通道缩放映射,使得分割效果得到少量的提升,但这个提升明显不如全连接CRF(两者可以兼容使用)。
论文:《DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs (CVPR2016)》
主要贡献:对同一张特征图使用不同dilation的空洞卷积,并用多孔空间金字塔池化下采样到同一尺寸,以此融合多种感受野的特征信息。
思路很简单,其实是何恺明的空间金字塔池化的一个演化。关于空间金字塔池化(SPP)的内容可以回顾《漫谈池化层 - 空间金字塔池化 | Hey~YaHei!》,此处不再赘述。
(图中rate指的就是我们常说的dilation,也即对输入特征图两个采样点的间隔,如标准卷积dilation=1)
论文:《Rethinking Atrous Convolution for Semantic Image Segmentation (CVPR2017)》
主要贡献:
具体pytorch实现可以参考yassouali/pytorch_segmentation/models/deeplabv3_plus.py#L249-L297
论文:《Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation (ECCV2018)》
主要贡献:
(原文以使用了深度可分离卷积的Xception为backbone)
如图,作者发现经过空间金字塔池化之后直接做一次八倍上采样过于粗暴,于是借鉴了编码器-解码器架构的思路,在空间金字塔池化后的特征图上采样四倍之后,与从中间层抽离出特征相融合,再做一次四倍上采样(注意编码器部分比原始模型要多一个阶段,所以与原输入相比编码器的输出实际上下采样了十六倍而非八倍,换句话v3+加深了网络)
可参考《漫谈卷积层 - 高效卷积 | Hey~YaHei!》,此处不再赘述。
参考:《图像分割领域常见的loss fuction有哪一些? | 知乎, 小锋子Shawn》
工具:mdbloice/Augmentor 支持在变换原图的时候同步操作ground truth,很方便