用MPSNN运行编译后的Core ML模型

Core ML已经出2.0了,性能进一步增强,但是之前就有的一个老问题还是存在:不支持任意大小图像的输入输出。虽然对于图像分类、语义分割来说问题不大,但是对图片滤镜类的卷积神经网络就不太友好了。MPSNN基于Metal Performance Shaders,作为iOS平台另一个机器学习框架,虽然学习曲线相对较陡,但是实际使用比Core ML灵活得多。除了支持任意图像大小外,还能动态更改模型,甚至可以将各种非机器学习的计算内核(Computing Kernel)也组合起来用。

iPhone XS/XS Max和新iPad Pro搭载了内置神经网络加速器的A12芯片。然而,waifu2x-ios在这块“性能怪兽”上似乎有点水土不服,输出图片会有明显的色阶(见此issue)。手上并没有设备进行调试,所以只能靠推测认为是因为神经引擎降精度运行导致的。即使有设备调试,也就能看看GPU Frame,多半也没有什么好解决办法,因为Core ML内部就是个黑盒,外面无从得知里面到底怎么运行的,更不可能通过一些手段介入。所以,是时候考虑增加MPSNN作为后端了。

Core ML的数据存储

作为更加底层的框架,MPSNN没有定义任何数据格式。神经网络中的weights和bias只接受按照顺序存储的32位浮点数指针。考虑到Core ML已经保存有包括网络结构在内的全部数据,完全有可能直接通过程序读取mlmodel中的数据构建MPSNN网络。这样一来既有了现成的数据格式,又可以节省另外存储模型的空间,还能同时兼容现有的Core ML模式。

经过搜索找到一篇详细讲解Core ML编译后格式的博客:Reverse Engineering Core ML

简单来说,Core ML在编译后会把模型的数据保存在以.mlmodelc结尾的文件夹中(这很Apple)。其中model.espresso.shape存储网络每一层的形状,实际上应该是被用来限制Core ML中feature的大小。model.espresso.net存储网络结构,为JSON格式。model.espresso.weights以32位浮点数的形式存储所有的weights和bias,数据没有任何加密。

通过mlmodelc构建MPSNN

重建网络

JSON文件中每个对象代表网络的一层,topbottom字段分别表示上一层和下一层的feature名称(和节点名称name不一样)。通过这些信息可以建立一个有向无环图,只要知道输入feature和输出feature的名称即可通过遍历这个图构建出路径。值得注意的是激活层(activation)和卷积层(convolution)是分开的,而在MPSNN的API中这两层需要被放到同一个节点(node)中。对于waifu2x这种只有一条路径的网络,构建起来很简单,甚至不需要考虑一个feature需要用两次的情况。

读取weights和bias

bias很简单,通过blob编号找到位置,按顺序复制就完事了。而weights相对来说就有点麻烦,因为Core ML和MPSNN对weights顺序的定义不一样。Core ML存储顺序为[outputFeatureChannels inputFeatureChannels kernelWidth kernelHeight],而MPSNN存储顺序为[outputFeatureChannels kernelHeight kernelWidth inputFeatureChannels]。而在反卷积核中,MPSNN的格式不变,Core ML的格式变成了[inputFeatureChannels outputFeatureChannels kernelWidth kernelHeight]。这种不一致导致必须在将weights传给MPSNN前将格式进行转换。

进行模型预测

由于完全运行在GPU上,MPSNN必须在真机上才可以运行。刚开始网络输出一片空白,通过抓取GPU Frame得知是因为alpha通道因为不明原因而全是0。这时MPSNN的优势就体现出来了,直接在网络输出的MPSImage后面增加一个kernel将alpha填1,输出就正常了。

代码实现: https://github.com/imxieyi/CoreML-MPS

文章目录
  1. 1. Core ML的数据存储
  2. 2. 通过mlmodelc构建MPSNN
    1. 2.1. 重建网络
    2. 2.2. 读取weights和bias
  3. 3. 进行模型预测
|