我挖了一个坑,然而没有时间填上了

五一劳动节当天,我看到 @sakamoto-poteko 坂本番茄酱前辈在他的知乎专栏里发了一个公告。我在一段时间前知道了他要从计算机转行到医生,如今他在公告里说他专栏里的『双倍的老婆』系列因此停更了。专栏的第一篇文章里有这么一句话,让我觉得特别有情怀:

为压片组租用天河二号铺平道路

因此我打算接过来自己挖这个坑。这两个星期刚刚挖出了一个坑位准备开始填,现实生活中又有了新的安排,我不得不停止填坑。挖坑过程我只做了一点渺小的工作,希望下文的内容对后面打算填坑的人有一点点帮助:

===============正文的分割线===============

一、waifu2x是什么,干什么,怎么用

原『双倍的老婆』系列的第一篇文章对waifu2x的原理有一个很抽象简单的介绍。简单说来,waifu2x通过一个七层神经网络,训练出了一组参数(一组卷积核),converter可以利用这一组参数来对其他动漫风格的图片进行放大和降噪,并且有非常良好的效果。这个网页里有日本人收集并维持更新的waifu2x各派生版本的情况,主要是在不同平台上的不同实现。如果你不打算训练自己的参数,只需要converter,那么对于Linux用户,可以从GitHub上下载Takashi Nakamura的版本,这个版本支持AVX、FMA、OpenCL和CUDA,应该是现在速度最快的版本(然而我不知道在没有CUDA支持的机器上能否编译);在Windows上也有预编译好的版本,我没用过,也就不去搜了。国内有些字幕组已经开始用waifu2x对某些人类文明四散的尘埃进行处理了。

二、waifu2x-converter的工作原理简介

『双倍的老婆』系列的第一篇文章 里有一些介绍,不过像我这种没有接触过CNN的人其实不知道到底是怎么弄的。在我自己挖过坑以后,我觉得waifu2x-converter的工作原理大概可以概括为如下:

首先将图片用bicubic算法upscale到目标尺寸上,然后将图像分割为若干个区域(可能重叠),每一次处理一小块区域。我们只讨论每次处理的这一小块区域(矩阵)。处理的流程如下(用语可能有些不规范,能理解就好):
1.将输入的目标矩阵用32个卷积核进行卷积处理,得到32个新的矩阵;
2.将上一步得到的32个矩阵,用32×32个卷积核进行处理,再次得到32个新的矩阵;
3.将上一步得到的32个矩阵,用32×64个卷积核进行处理,得到64个新的矩阵;
4.将上一步得到的64个矩阵,用64×64个卷积核进行处理,再次得到64个新的矩阵;
5.将上一步得到的64个矩阵,用64×128个卷积核进行处理,得到128个新的矩阵;
6.将上一步得到的128个矩阵,用128×128个卷积核进行处理,再次得到128个新的矩阵;
7.将上一步得到的128个矩阵,用128×1个卷积核进行处理,最后得到1个目标输出矩阵。

不要问我为什么是这个过程,我没有读过原来的论文,现在自然也没有时间再去读原来的论文了,你们看我是不是很不要脸啊哈哈哈哈。所以说,其实converter的工作流程不是很复杂,就是做了7轮卷积。但是上面的过程漏了一点,每一轮的卷积和它的输入输出的矩阵数是什么关系呢?

我们约定(其实是代码里的变量名)每一轮有nInputPlanes个矩阵输入,要输出nOututPlanes个矩阵。那么,我们需要nInputPlanes * nOutputPlanes个卷积核。这些卷积核都是已经训练出来的3×3的矩阵。每一个输入的矩阵都要用不同的卷积核进行卷积,然后将卷积结果加到每一个输出的矩阵上。处理流程的伪代码如下:

 

我测了一下,对于原始的程序,处理时调用OpenCV的filter2D函数,这个函数对于输入矩阵的边缘处理,相当于将下方左侧的n×n矩阵填充成右侧的(n+2)x(n+2)矩阵,然后对填充后中间的n×n矩阵做卷积:原来的四个角落复制到新的四个角落,原来的四条边复制到新的四条边。至于这个ReLU,其实就两步:首先,给矩阵里的每一个值都加上一个偏置修正量biases[opIndex],然后将矩阵里所有小于0的元素乘以0.1。

好了,converter的所有工作流程已经讲清楚了。

三、挖坑的原点

我选用的基线版本是WL-Amigo的waifu2x-converter-cpp,Takashi Nakamura的版本也是从此版本上发展出来的。之所以选择这个版本,是因为它是C++写的(别的语言我没学过,而且Xeon Phi只支持C/C++/Fortran),最简单(太复杂的我看不懂更加改不动)。

我用来开发的机器是实验室里的一台服务器,CentOS 7.1,Xeon E5-2620v2 * 2 @ 2.1GHz(12C24T),6*16GB DDR3 1600MHz,1 * Xeon Phi 31S1P和1 * Tesla K40c,ICC 2016。CPU理论单精度峰值403.2G,K40c理论单精度峰值4.3T,31S1P理论单精度峰值约2T(offload模式下的sgemm实测只有1500G),内存带宽应该是76.8GB/s。

要编译这个基线版本的程序,我们需要先编译一个OpenCV 3.0。OpenCV 3.0的简要编译步骤在我的waifu2x-converter-XeonPhi的README里有。编译完以后我没有用pkg-config,因为如果以后要在天河二号或者其他超算上使用,你是没有办法用sudo权限的。好在不需要pkg-config也可以用。一开始我编译OpenCV直接用了ICC,但是发现OpenCV需要下载Intel的IPP作为第三方库,这和我机器上的ICC自带的有冲突,而且我不知道要如何处理,就干脆用GCC来编译了。基线版本的程序没有Makefile,我自己写了一个凑合对付。同时,我们还需要修改原来的modelHandler.cpp的第153、154行的错误,将其改为这样。如此编译出来的程序,在服务器上用CPU(因为这个程序禁用了OpenCV的OpenCL)跑出来的峰值速度大概是30GFlops多一点。于此形成对比的是,Takashi Nakamura的纯CPU版本峰值速度是320G,GPU峰值速度是820G左右。

四、坑里的模样

现在让我们来看看这份代码在哪里进行计算。基线版本程序的计算入口在modelHandler.cpp的Model::filter(),这个函数负责清空输出矩阵组,划分任务,并且启动多个线程执行计算核心函数 Model::filterWorker()。调用Model::filter()的时候,参数是输入矩阵组和输出矩阵组的引用,nInputPlanes、nOutputPlanes、nInputPlanes * nOutputPlanes个卷积核和偏置修正值都保存在Model这个对象的私有成员里。所以,我一开始选择在Model::filter()里Hack掉Model::filterWorker(),替换成自己的卷积实现,然后再一点一点进行修改。

为了复用Xeon Phi加速卡上申请的内存空间,我修改了一下流程。调用Model::filter()的是convertRoutine.cpp里的convertWithModelsBasic(),因此我在这个函数内调整了一下调用顺序:
1.统计出nInputPlanes * nOutputPlanes的最大值(其实可以直接写死是128×128),当前处理区域的大小,将其送入initLocalMem(),建立myConvKernel.cpp内的本地全局数组,并在Xeon Phi上一次过分配完毕所有空间;
2.按照原来的顺序依次调用Model::filter(),执行以下操作:
(1)调用copyInMatrices(),将这一轮的卷积核数组和偏置修正值数组上传到Xeon Phi,并且在Xeon Phi上将上一轮的outputPlanes重新填充边缘变成这一轮的inputPlanes;
(2)调用myConvKernel()进行这一轮的计算;
3.调用copyOutResults()将计算结果从Xeon Phi上下载下来,并释放Xeon Phi加速卡上申请的所有空间,下载下来的计算结果写回去原来程序的cv::Mat类型的outputPlane中。

这是一个很Naive的实现,因为只是单纯地做了一个『移植』,而没有任何的优化。经测试,我的纯CPU版本的峰值速度只有120G左右,总处理速度约100G;Xeon Phi版本峰值速度只有250G,总处理速度210G左右。显然,这距离他们各自的峰值速度还差很远,起码CPU上要达到峰值速度的三分之二以上、Xeon Phi上要达到峰值速度的25%(GPU版的在其他GPU上也没有超过30%的)左右才能算是可以使用的状态。

五、一些还没用得上的填坑工具

要利用好机器,还是要从Cache利用率、向量化和并行方法这三个角度去下工夫。并行方法我是基于行的并行,即每个线程处理一个矩阵的若干行,这是一个比较容易的方法,也是Takashi Nakamura使用的。但是我没有时间去重新设计程序的计算顺利来提高Cache利用率了。我给Takashi Nakamura发过邮件请教过这一问题,他表示他的文档有些老旧,他对数据的重新打包和向量化可以参考这个这个这个,最后一个在waifu2x-converter-cpp的Network Graph上还可以看到接连的好几个关于packed input/output的更新。至于他现在的处理方法,他给出了如下的伪代码:

看上去他对数据的打包还是比较复杂的。希望有人能弄出一个更简单的方法吧。

顺带多说一句,由于有ICC,所以我基本不需要手写向量化指令,这是ICC最好的一个地方。

===============正文的结束线===============

再次重申,我所实现的waifu2x-converter-XeonPhi只是一个原型,未经过充足的优化和测试,距离真正可以拉到天河二号上面去压片还有很长的距离。在此向坂本番茄酱前辈的脑洞致敬,并祝他成为一个优秀的麻醉大夫,麻倒一片小萝莉

最后强行贴一张图。这是某waifu的CG原图(720P):

waifu0

天河二号同款的Xeon Phi 31S1P上的8×28个框框神秘力量加成的结果(1440P):

waifu2x-MIC

我挖了一个坑,然而没有时间填上了》上有5条评论

发表评论

电子邮件地址不会被公开。 必填项已用*标注