本文已参与「新人创作礼」活动，一起开启掘金创作之路。

Portaitnet: Real-time portrait segmentation network for mobile device

概述：为了实现对图片进行人像分割，但现有网络对具有模糊的边界和复杂的光照条件的图片无法实现精细分割，而且这些网络结构很大，不适合移动端实时快速的人像分割；作者为了解决上述问题，提出了一个实时人像分割模型，PortraitNet，它是一个基于轻量级的U型架构，除了对分割结果计算交叉熵损失，还用到两个辅助损失，边界损失和一致性损失，能够提高边界像素的分割精度和增强了复杂光照环境下的鲁棒性，它可以有效地在移动设备上运行。

1. 预备知识

1.1 MobilenNet V2

MobilenNet V2:Inverted Residuals and Linear Bottlenecks: Mobile Networks for Classiﬁcation, Detection and Segmentation 从MobilenNet V2的题目就能得知，它是一种用于分类，目标检测和语义分割的移动网络，它主要用到的结构是Inverted Residuals 和 Linear Bottleneck，就是一个具有线性瓶颈的倒置残差模块，首先该模块将输入的低维压缩表示拓展到高维而且用轻量级深度卷积进行过滤，然后用线性卷积将特征投影回低维表示。

Inverted residual block: expand - transfer - reduce （中间宽，两头窄）因为Depthwise卷积计算量很少，所以需要先用 1x1 卷积提高通道数，再用 Depthwise 3x3 卷积，最后再用 1x1 卷积降低通道数。两端的通道数都很小，而且depthwise卷积计算量也不大，所以这种基于这种结构的模型比较小，能够用在移动端设备上。
Linear Bottleneck 就是最后 1x1 卷积后不用 ReLU，而是linear，目的是防止ReLU破坏特征。
Inverted residual block + Linear Bottleneck 针对stride=1 和stride=2，在block上有稍微不同，主要是为了与shortcut的维度匹配，因此， stride=2时，不采用shortcut。具体如下图：

作者将 Mobilenet V2 作为主干特征提取网络，输入的是图片，输出的是每一层的feature。

1.2 U-Net

U-Net: Convolutional Networks for Biomedical Image Segmentation Unet是解决图像分割的一种常见网络结构。 Unet可以分为三个部分，如下图所示：第一部分是主干特征提取部分，我们可以利用主干部分获得一个又一个的特征层，Unet的主干特征提取部分与VGG相似，为卷积和最大池化的堆叠。利用主干特征提取部分我们可以获得五个初步有效特征层，在第二步中，我们会利用这五个有效特征层可以进行特征融合。第一部分对应下图中左边部分，也就是从输入到遇到第一个绿色箭头的部分。第二部分是加强特征提取部分，我们可以利用主干部分获取到的五个初步有效特征层进行上采样，并且进行特征融合，获得一个最终的，融合了所有特征的有效特征层。第三部分是预测部分，我们会利用最终获得的最后一个有效特征层对每一个特征点进行分类，相当于对每一个像素点进行分类。

2. 算法原理

2.1 模型结构

整个网络采用了经典的Unet结构，如下图所示，绿色块代表编码器模块，括号中的数字代表下采样速率，为了提高速度，backbone是Mobilenet V2；黄色和紫色块代表解码器模块，与Unet不同的是，编码器使用残差结构的模块D-Block，D-Block块有两个分支，一个分支包含两个深度可分离卷积以提高效率，另一个包含单一的 1x1 卷积来调整通道的数量。在模型的最后采用了多头输出结构，输出 mask的同时，输出图片分离的边界信息，将这两个输出用于训练。

2.2 损失函数

除了对输出的mask与ground truth的mask计算交叉熵损失函数。还有两个辅助损失函数，一个是能够提高边界像素的分割精度的边界损失，另一个增强复杂光照环境下的鲁棒性的一致性损失。

边界损失因为人像分割对分割边界很敏感，所以网络需要生成更清晰的边界，用来支持背景模糊或替换等应用。做法是在训练时最后一层，增加了一个卷积层用于预测边界，计算损失函数，而在测试时不用预测边界。因为边界占据图像很小的部分，为了避免极度的样本不均衡，所以用的是focal损失函数，公式如下。

$y_i$ 代表像素的ground truth标签， $p_i$ 代表像素i 的预测概率。边界损失可以提高模型对人像边界的敏感性，进而提高分割精度.

一致性损失人像图片在不同光照条件下，得到亮度不同但内容相同的图片，这些图片的标签是一样的，但网络可能因为不同亮度下的像素值而得到不同的分割预测，为了避免这种情况，增强模型的鲁棒性，采用一致性损失。它采用软标签和KL散度进行训练。具体做法是对于原始图像，我们先用变形增强（旋转，缩放）的方法生成图像，然后对图像进行纹理增强（改变亮度，对比度）生成，纹理增强不会改变图像的形状，因此图像和的分割是相同的。假设图像的网络输出为热图，图像的输出为热图，而和内容是相同的，所以直接和作为硬标签的ground truth计算交叉熵损失。由于具有更多信息的软标签有利于训练模型，所以将质量更好的作为软标签，用于和计算KL散度损失，即一致性损失，公式如下。

$L_m$ '是网络预测的B , B'与真实标签的交叉熵损失。

3. 算法实现

3.1 实验环境

框架：pytorch 显存：8G 运行内存：64G 代码链接：github.com/Bailey-24/P…

3.2 数据集

用的两个数据集，并用数据增强来补充原始数据集，提高训练模型的通用性，从而获得更好的分割结果。第一个数据集是EG1800，使用1447张图片进行训练，289张进行验证。第二个数据集是Supervise-Portrait，使用1858张进行训练，400张进行验证。与 EG1800 相比， Supervise-Portrait 中的人像图像具有更复杂的背景和严重的遮挡。数据增强包括变形增强和纹理增强。变形增强包括随机水平翻转，旋转，缩放和平移。纹理增强包括随机噪声，图像模糊，颜色，亮度，对比度，锐度变化。以下为对原始数据集进行增强和计算人像边界的核心代码。

3.4 实验结果

在EG1800数据集上跑了2000个epoch后的精度为0.9672。下面的图为我的模型在EG1800数据集对人像进行前后景分割的结果，左边为原图，右边为结果。

在Supervisely face数据集上跑了2000个epoch后的精度为0.9326。下面的图为我的模型在Supervisely face数据集对人像进行前后景分割的结果，左边为原图，右边为结果。

4. 心得体会

最大的感触就是我们应该学会灵活运用工具。 Jittor是一个性能好，速度快的深度学习框架，但我用下来觉得有两个问题，首先它的API相比其他框架要少，所以一些损失函数的API要自己写，比如论文用到的focal loss和数据增强的函数，写完之后调试时，报错信息给的很模糊，就很难找到错误原因，因此要经常去论坛问jittor的作者。其次，jittor编译完运行速度很快，但我还没抓清楚它使用的显存和内存分配机制，所以由于硬件问题跑不起来。我有好几次编译代码就花了半小时，结果编译完跑起来就报错，后面了解到原因是计图的算子创建和执行是分开的，所以它需要更多显存来存储，我电脑不够显存，服务器要15号才能分配给我，等有服务器之后再用jittor试一下。总结一下，Jittor有丰富的图像学相关的深度学习库，比如GAN模型库，三维点云库，可微渲染库，这些库的API比较丰富，但其他任务或工作还是尽量使用比较成熟的框架来完成，可以节约重复造轮子的时间。等拿到服务器之后，再测试jittor相比其他框架的速度和性能。

5. 参考文献

PortraitNet: Real-time portrait segmentation network for mobile device论文解读和代码实践
pytorch_PortraitNet
PortraitNet论文复现以及其中的数据增强pp飞浆
jittor示例4：语义分割之 DeepLabV3+
pytorch模型辅助转换jittor工具教程
U-Net: Convolutional Networks for Biomedical Image Segmentation.
轻量级网络--MobileNetV2论文解读
PortraitNet-Paddle/myTest/VideoTest
jittor与torch之间的差异

Portaitnet: Real-time portrait segmentation network for mobile device论文笔记