MATLAB 一些向量化思路和例子

Posted on 2020-05-21 Last updated on 2020-05-23

Digital Image Processing的处理经常需要密集的计算，效率低的代码处理一张500x500的图片都可以计算30秒。~~为了方便调参~~，~~同时身为超算比赛参赛者，职业病犯了~~，写一个速度快的代码是很有必要的。本文的code和benchmark来自我写的DIP实验报告。

Convolution / Correlation

卷积 (Convolution)是DIP和CNN常用的一种操作，而Correlation是一种和卷积差不多的操作，把卷积核转置两下做卷积就是Correlation。Correlation在DIP可以用来做模糊和边缘检测。下面我给大家表演徒手写向量化Correlation。

常规思路

假设padImgArr输入的图片，mask就是类似卷积核的一个东西（mask转置两下做Correlation就是Convolution）。正常思路当然是两层for循环，遍历图片的每一个像素，然后取像素临近的点，和mask来一发点乘。菜一点的人会再来两层for循环，去逐个计算像素值乘mask值的结果，稍微觉得有点不对劲的人，很容易想到对padImgArr切片做点乘再求和，这种想法的代码实现如下。

function outImgArr = correlation2(padImgArr, mask)
    imgSize = size( padImgArr );
    outImgArr = zeros( imgSize );
    H = imgSize(1); W = imgSize(2);
    for i = 2:H-1
        for j = 2:W-1
            outImgArr(i, j) = sum( padImgArr( i-1:i+1, j-1:j+1 ).*mask, "all" );
        end
    end
end

优化方法

如果你是一个对解释型语言的循环有PTSD的人，很快你会开始考虑如何减少循环次数。

换一个思路，对mask的元素做一次遍历，而不是遍历每个像素，如果是3x3的mask，512x512的padImgArr，那么循环次数就从512x512=262144次降低到3*3=9次。很容易~~（并不是）~~就可以发现，在整个计算过程中，mask左上角的橙色元素只会和图片坐上角橙色框的像素相乘，其数组乘积是最终结果的一部分。完善这个思路~~（完整思路看图片意会去）~~，可以写出下面的代码。

function outImgArr = correlation(padImgArr, mask)
    imgSize = size( padImgArr );
    outImgArr = zeros( imgSize );
    H = imgSize(1); W = imgSize(2);
    for i = -1:1
        for j = -1:1
            if mask(i+2, j+2)
                outImgArr(2:H-1, 2:W-1) =  outImgArr(2:H-1, 2:W-1) + padImgArr( (2+i):(H-1+i), (2+j):(H-1+j) )*mask(i+2, j+2);
            end
        end
    end
end

if mask(i+2, j+2)这一句是判断mask的元素是否为0，为0就跳过计算，反正0乘什么东西都是0，这对稀疏的mask有一定的加速作用。

最终优化的结果如下

这个优化产生了10倍左右的加速，除了因为避免解释型语言的循环开销，还有可能是因为每一次向量计算的计算规模增大，SIMD非对齐内存访问的次数减少。

Max-pooling / Max Filter

CNN的最大池化 (Max-pooling)和DIP的Max Filter在工作方式上差不多，只不过前者是用来缩小计算规模，后者经常拿来消除图片的某种噪声。最大池化通常配置是跨步为2，池化核为2x2，而Max Filter用CNN的黑话来说就是通常配置为跨步为1，池化核为3x3，5x5等。接下来以DIP的Max Filter为例，向量化MATLAB的代码。

常规思路

常规思路就是常规思路，每个像素点遍历一下，该干什么就干什么，就像上面的常规思路一样搞，代码简单朴素又原始。

function outImgArr = reduceSAP2(inImgArr, nSize)
    edgeW = floor( nSize/2 );
    padImgArr = padImg( inImgArr, [edgeW edgeW], 0, "center");
    outImgArr = zeros(size(padImgArr));
    inImgSize = size( inImgArr );
    inH = inImgSize(1); inW = inImgSize(2);
    for i = edgeW+1:inH+edgeW
        for j = edgeW+1:inW+edgeW
            outImgArr(i, j) = max( padImgArr(i-edgeW:i+edgeW, j-edgeW:j+edgeW), [], 'all');
        end
    end
    outImgArr = uint8( outImgArr(edgeW+1:inH+edgeW, edgeW+1:inW+edgeW) );
end

这里的padImg函数是我写的用来给图像pad一圈0的。如果是在3x3范围内取最大值，那么nSize就为3，edgeW就为1，意味着pad了一圈0。如果edgeW为2，就pad两圈0。

优化方法

如果你是PTSD和OCD患者，就往下看吧。

一个简单的想法是把计算每一个像素需要用到的那几个数据全部给拎出来存着，很快~~（并没有）~~就联想起了之前的优化思路，对原始图片疯狂移位。这个时候建立一个三维数组，对原始图片疯狂移位，让临近的像素蹭到中心像素位置上，然后把临近的像素值保存到不同通道上（借用一下CNN的概念，看图片很容易理解），相当于保存了一堆错位的图片，最后把这一堆叠在一起的图片啪叽一下拍成一张图片就好了。顺着这个感觉可以写出来以下的实现。

function outImgArr = reduceSAP(inImgArr, nSize)
    inImgSize = size( inImgArr );
    inH = inImgSize(1); inW = inImgSize(2);
    
    edgeW = floor( nSize/2 );
    padImgSize = inImgSize + 2*edgeW;
    stackImgArr = repmat(0, [padImgSize nSize^2]);
    
    for i = -edgeW:edgeW
        for j = -edgeW:edgeW
            stackImgArr(edgeW+1+i:inH+edgeW+i, edgeW+1+j:inW+edgeW+j, nSize*(i+edgeW)+(j+edgeW)+1) = inImgArr;            
        end
    end
    
    stackImgArr = stackImgArr(edgeW+1:inH+edgeW, edgeW+1:inW+edgeW, :);
    outImgArr = uint8( max(stackImgArr, [], 3) );
end

注意这里的max函数指定了参数3，意思是把第三维上所有的数据规约，变成单个值，这样就可以让三维数组降维成二维数组。

优化的效果也很还行，也就那么100倍的性能提升。

牢骚

一开始让我写MATLAB，我是拒绝的。你看MATLAB的代码又臭又丑（~~看到了其他电子系同学的代码后~~）。后来发现随便一个MATLAB的向量操作，就有多线程优化，AVX优化，执行效率还高，同时Python的多线程性能还是那么感人，C语言起手先来句#pragma omp parallel for，接上一个for循环，再来两句代码，才等效一句MATLAB代码，而且说不定还引入了两个bug...MATLAB真香。