移动端人像分割开发

作者 | Video++极链科技AI实验室谢锦滨

整理 | 包包

个人对移动端神经网络开发一直饶有兴致。去年腾讯开源了NCNN框架之后，一直都在关注。近期尝试将分割网络移植到NCNN，能够在手机端实现一些有趣的应用，本文就几个技术话题作相关介绍。

神经网络选择

目前做segmentation常用的一些神经网络网络有如下几个可以选择:

MaskRCNN

FCN

UNET

SegNet

Tiramisu

在移动端做人像分割有两大优势，首先是隐私，其次是可以做到实时，能够创造更多玩法。因为UNET模型比较简单，干脆就从这个入手。下面是UNET网络结构：

首先我采用了基于keras的版本: https://github.com/TianzhongSong/Person-Segmentation-Keras，训了一个基本模型，大小为39M, iphone X上15秒处理一帧。明显这个速度太慢，需要进行改造。

移动端Inference框架

经过调研，粗略比较了几个神经网络框架:

其中使用难易程度，主要跟我个人习惯有关。NCNN框架比较好，代码不多，而且兼容iOS和安卓（台式机以及嵌入式环境同样支持），同时底层计算采用汇编做了优化。NCNN只实现神经网络的forward部分，没有反向传播，所以训练仍旧依赖其他开源框架，现在几大框架都遵守ONNX协议，理论上各种框架模型之间互相转换并不存在什么问题，工具也都是开源的。

不过keras没办法直接转成ncnn模型，研究过通过onnx模型做中间跳板，采用了一些开源的转换工具，也是一堆问题。NCNN支持几个神经网络训练框架:caffe/mxnet/pytorch，在ncnn的github有一篇issue里nihui推荐采用MXNET，因此MXNET也成为了我的首选。其他框架往NCNN转换工具:

NCNN转换Tensorflow模型有问题； Caffe没有Pytorch和MXNET好用; 最终在MXNet和Pytorch之间选择了MXNet。

人像数据集

https://github.com/lemondan/HumanParsing-Dataset

https://github.com/ZhaoJ9014/Multi-Human-Parsing_MHP

COCO人像数据集 – 加入后效果质的飞跃

ADE20K

网上找了上面几个数据集，抽取出人像部分，采用基本的flip/crop/rotate操作做了扩充，得到228423张训练样本，另外凑了9064张验证样本。

模型转换（MXNET->NCNN）

MXNET的UNET版本并没有现成可用的合适版本。参照其他版本的UNET，自己coding完成一个版本。代码请参考: https://github.com/xuduo35/unet_mxnet2ncnn。

在这个基础上训练完成，用来测试ncnn转换基本可用。这里提一下转换过程遇到的一些问题和解决方案。

一个是调用ncnn extract函数会crash，经过调查，发现mxnet2ncnn工具也有bug，blob个数算错，其次是input层one_blob_only标志我的理解应该是false，不知道什么原因转换过来的模型这边是true，导致forward_layer函数里面bottoms变量访问异常。后来一层层extract出来打印输出的channel/width/height调查后又发现，自己代码里unet.py里的name为pool5写成了pool4，前面的crash跟这个致命错误有关系也有直接关联。

第二个问题是转成ncnn后的预测结果死活不对。只能一层层去检查，写了几个简单的工具可以打印中间隐藏层的结果（代码: https://github.com/xuduo35/unet_mxnet2ncnn/check.py）。在这个基础之上，发现是第一次反卷积就出了问题（mxnet神经网络trans_conv6的输出）。结果完全不一致，按个人理解，反卷积算法会出问题的可能性基本为0，所以把mxnet这一层的权重值打印了出来。再在mxnet2ncnn的代码里把对应的参数打印，最后发现是num_group出了问题，简单处理就是把mxnet2ncnn.cpp里的反卷积num_group固定为1，终于解决问题。得到正确的输出结果：

中间还遇到一些ncnn和mxnet之间图像格式之类的转换问题，特别是浮点数的处理，就不啰嗦了。另外，调试过程发现，ncnn的中间层输出和mxnet的输出不是完全一致，可能是有一些参数或者运算细节问题，不影响最后mask结果，暂时忽略。

几个问题

到目前为止还存在几个问题，1. 模型比较大；2. 单帧处理需要15秒左右的时间（Mac Pro笔记本，ncnn没有使用openmp的情况）；3. 得到的mask结果不是特别理想。针对这三个问题，对网络结构进行调整。

1. 模型比较大

采取将网络卷积核数量减少4倍的方式，模型大小下降到2M，粗略用图片测试，效果也还可以。同时把之前用0值填充图片的方式，改成用边界值填充，因为测试的时候发现之前的方式总在填充的边界往往会出现检测错误。

2. 单帧处理需要15秒左右的时间

按照第一步处理之后，基本上一张图片只要1秒钟就处理完成，。在手机上由于NCNN做了优化，经过测试速度是Mac Pro的好几倍。

3. 得到的mask结果不是特别理想

在权衡模型大小和准确率的基础上修改UNET网络结构，具体不再赘述。

最终结果