numpy中二维的跨步卷积

Strided convolution of 2D in numpy

我尝试使用 for 循环实现二维数组的跨步卷积,即

arr = np.array([[2,3,7,4,6,2,9],
                [6,6,9,8,7,4,3],
                [3,4,8,3,8,9,7],
                [7,8,3,6,6,3,4],
                [4,2,1,8,3,4,6],
                [3,2,4,1,9,8,3],
                [0,1,3,9,2,1,4]])

arr2 = np.array([[3,4,4],
                 [1,0,2],
                 [-1,0,3]])

def stride_conv(arr1,arr2,s,p):
    beg = 0
    end = arr2.shape[0]
    final = []
    for i in range(0,arr1.shape[0]-1,s):
        k = []
        for j in range(0,arr1.shape[0]-1,s):
            k.append(np.sum(arr1[beg+i : end+i, beg+j:end+j] * (arr2)))
        final.append(k)

    return np.array(final)

stride_conv(arr,arr2,2,0)

这导致 3*3 数组:

array([[ 91, 100,  88],
       [ 69,  91, 117],
       [ 44,  72,  74]])

是否有 numpy 函数或 scipy 函数来做同样的事情?我的方法不是很好。我怎样才能矢量化这个?

忽略填充参数和尾随 windows,它们没有足够的长度与第二个数组进行卷积,这是 np.lib.stride_tricks.as_strided -

的一种方法
def strided4D(arr,arr2,s):
    strided = np.lib.stride_tricks.as_strided
    s0,s1 = arr.strides
    m1,n1 = arr.shape
    m2,n2 = arr2.shape    
    out_shp = (1+(m1-m2)//s, m2, 1+(n1-n2)//s, n2)
    return strided(arr, shape=out_shp, strides=(s*s0,s*s1,s0,s1))

def stride_conv_strided(arr,arr2,s):
    arr4D = strided4D(arr,arr2,s=s)
    return np.tensordot(arr4D, arr2, axes=((2,3),(0,1)))

或者,我们可以使用内置的 scikit-image view_as_windows 来优雅地 windows ,像这样 -

from skimage.util.shape import view_as_windows

def strided4D_v2(arr,arr2,s):
    return view_as_windows(arr, arr2.shape, step=s)

这是一个基于 O(N^d (log N)^d) fft 的方法。这个想法是将两个操作数在所有偏移量模步幅处分割成步幅间隔的网格,在相应偏移量的网格之间进行常规的 fft 卷积,然后逐点求和结果。它有点重索引,但恐怕无能为力:

import numpy as np
from numpy.fft import fftn, ifftn

def strided_conv_2d(x, y, strides):
    s, t = strides
    # consensus dtype
    cdt = (x[0, 0, ...] + y[0, 0, ...]).dtype
    xi, xj = x.shape
    yi, yj = y.shape
    # round up modulo strides
    xk, xl, yk, yl = map(lambda a, b: -a//b * -b, (xi,xj,yi,yj), (s,t,s,t))
    # zero pad to avoid circular convolution
    xp, yp = (np.zeros((xk+yk, xl+yl), dtype=cdt) for i in range(2))
    xp[:xi, :xj] = x
    yp[:yi, :yj] = y
    # fold out strides
    xp = xp.reshape((xk+yk)//s, s, (xl+yl)//t, t)
    yp = yp.reshape((xk+yk)//s, s, (xl+yl)//t, t)
    # do conventional fft convolution
    xf = fftn(xp, axes=(0, 2))
    yf = fftn(yp, axes=(0, 2))
    result = ifftn(xf * yf.conj(), axes=(0, 2)).sum(axis=(1, 3))
    # restore dtype
    if cdt in (int, np.int_, np.int64, np.int32):
        result = result.real.round()
    return result.astype(cdt)

arr = np.array([[2,3,7,4,6,2,9],
                [6,6,9,8,7,4,3],
                [3,4,8,3,8,9,7],
                [7,8,3,6,6,3,4],
                [4,2,1,8,3,4,6],
                [3,2,4,1,9,8,3],
                [0,1,3,9,2,1,4]])

arr2 = np.array([[3,4,4],
                 [1,0,2],
                 [-1,0,3]])

print(strided_conv_2d(arr, arr2, (2, 2)))

结果:

[[ 91 100  88  23   0  29]
 [ 69  91 117  19   0  38]
 [ 44  72  74  17   0  22]
 [ 16  53  26  12   0   0]
 [  0   0   0   0   0   0]
 [ 19  11  21  -9   0   6]]

我想我们可以做一个 "valid" fft 卷积,然后只挑选出跨步位置的结果,像这样:

def strideConv(arr,arr2,s):
    cc=scipy.signal.fftconvolve(arr,arr2[::-1,::-1],mode='valid')
    idx=(np.arange(0,cc.shape[1],s), np.arange(0,cc.shape[0],s))
    xidx,yidx=np.meshgrid(*idx)
    return cc[yidx,xidx]

这给出了与其他人的答案相同的结果。 但我想这只有在内核大小为奇数时才有效。

此外,为了与其他人保持一致,我在 arr2[::-1,::-1] 中翻转了内核,您可能希望根据上下文省略它。

更新:

我们目前有几种不同的方法可以单独使用 numpy 和 scipy 进行 2D 或 3D 卷积,我想做一些比较以了解哪种方法在不同大小的数据上更快。希望不要被当成off-topic.

方法一:FFT卷积(使用scipy.signal.fftconvolve):

def padArray(var,pad,method=1):
    if method==1:
        var_pad=numpy.zeros(tuple(2*pad+numpy.array(var.shape[:2]))+var.shape[2:])
        var_pad[pad:-pad,pad:-pad]=var
    else:
        var_pad=numpy.pad(var,([pad,pad],[pad,pad])+([0,0],)*(numpy.ndim(var)-2),
                mode='constant',constant_values=0)
    return var_pad

def conv3D(var,kernel,stride=1,pad=0,pad_method=1):
    '''3D convolution using scipy.signal.convolve.
    '''
    var_ndim=numpy.ndim(var)
    kernel_ndim=numpy.ndim(kernel)
    stride=int(stride)

    if var_ndim<2 or var_ndim>3 or kernel_ndim<2 or kernel_ndim>3:
        raise Exception("<var> and <kernel> dimension should be in 2 or 3.")

    if var_ndim==2 and kernel_ndim==3:
        raise Exception("<kernel> dimension > <var>.")

    if var_ndim==3 and kernel_ndim==2:
        kernel=numpy.repeat(kernel[:,:,None],var.shape[2],axis=2)

    if pad>0:
        var_pad=padArray(var,pad,pad_method)
    else:
        var_pad=var

    conv=fftconvolve(var_pad,kernel,mode='valid')

    if stride>1:
        conv=conv[::stride,::stride,...]

    return conv

方法 2:特殊转换(参见 this anwser):

def conv3D2(var,kernel,stride=1,pad=0):
    '''3D convolution by sub-matrix summing.
    '''
    var_ndim=numpy.ndim(var)
    ny,nx=var.shape[:2]
    ky,kx=kernel.shape[:2]

    result=0

    if pad>0:
        var_pad=padArray(var,pad,1)
    else:
        var_pad=var

    for ii in range(ky*kx):
        yi,xi=divmod(ii,kx)
        slabii=var_pad[yi:2*pad+ny-ky+yi+1:1, xi:2*pad+nx-kx+xi+1:1,...]*kernel[yi,xi]
        if var_ndim==3:
            slabii=slabii.sum(axis=-1)
        result+=slabii

    if stride>1:
        result=result[::stride,::stride,...]

    return result

方法 3:Strided-view conv,如 Divakar 所建议:

def asStride(arr,sub_shape,stride):
    '''Get a strided sub-matrices view of an ndarray.

    <arr>: ndarray of rank 2.
    <sub_shape>: tuple of length 2, window size: (ny, nx).
    <stride>: int, stride of windows.

    Return <subs>: strided window view.

    See also skimage.util.shape.view_as_windows()
    '''
    s0,s1=arr.strides[:2]
    m1,n1=arr.shape[:2]
    m2,n2=sub_shape[:2]

    view_shape=(1+(m1-m2)//stride,1+(n1-n2)//stride,m2,n2)+arr.shape[2:]
    strides=(stride*s0,stride*s1,s0,s1)+arr.strides[2:]
    subs=numpy.lib.stride_tricks.as_strided(arr,view_shape,strides=strides)

    return subs

def conv3D3(var,kernel,stride=1,pad=0):
    '''3D convolution by strided view.
    '''
    var_ndim=numpy.ndim(var)
    kernel_ndim=numpy.ndim(kernel)

    if var_ndim<2 or var_ndim>3 or kernel_ndim<2 or kernel_ndim>3:
        raise Exception("<var> and <kernel> dimension should be in 2 or 3.")

    if var_ndim==2 and kernel_ndim==3:
        raise Exception("<kernel> dimension > <var>.")

    if var_ndim==3 and kernel_ndim==2:
        kernel=numpy.repeat(kernel[:,:,None],var.shape[2],axis=2)

    if pad>0:
        var_pad=padArray(var,pad,1)
    else:
        var_pad=var

    view=asStride(var_pad,kernel.shape,stride)
    #return numpy.tensordot(aa,kernel,axes=((2,3),(0,1)))
    if numpy.ndim(kernel)==2:
        conv=numpy.sum(view*kernel,axis=(2,3))
    else:
        conv=numpy.sum(view*kernel,axis=(2,3,4))

    return conv

我做了3组对比:

  1. 对二维数据进行卷积,具有不同的输入大小和不同的内核大小,stride=1,pad=0。结果如下(颜色为卷积重复10次所用时间):

所以 "FFT conv" 通常是最快的。 "Special conv" 和 "Stride-view conv" 随着内核大小的增加变慢,但随着它接近输入数据的大小时再次减小。最后一个子图显示了最快的方法,所以紫色的大三角形表示 FFT 是赢家,但请注意左侧有一个细绿色柱(可能太小看不见,但它就在那里),表明 "Special conv" 对于非常小的内核(小于大约 5x5)具有优势。当内核大小接近输入时,"stride-view conv" 最快(见对角线)。

比较 2:3D 数据上的卷积。

设置:pad=0,stride=2,输入维度=nxnx5,kernel shape=fxfx5.

当内核大小在输入中间时,我跳过了 "Special Conv" 和 "Stride-view conv" 的计算。基本上 "Special Conv" 现在没有优势,而且 "Stride-view" 对于大小内核都比 FFT 快。

一个附加说明:当大小超过 350 时,我注意到 "Stride-view conv".

的内存使用量达到了相当大的峰值

比较 3:在 3D 数据上以更大的步幅进行卷积。

设置:pad=0,stride=5,输入维度=nxnx10,kernel shape=fxfx10.

这次省略了"Special Conv"。对于更大的区域,"Stride-view conv" 超过 FFT,最后的子图显示差异接近 100%。 可能是因为随着步幅的增加,FFT 方法会浪费更多的数字,因此 "stride-view" 对于小内核和大内核都会获得更多优势。

使用 signal.convolve2d from scipy 怎么样?

我的方法与 Jason 的方法类似,但使用索引。

def strideConv(arr, arr2, s):
    return signal.convolve2d(arr, arr2[::-1, ::-1], mode='valid')[::s, ::s]

请注意,必须反转内核。详情请见讨论 and here. Otherwise use signal.correlate2d.

示例:

 >>> strideConv(arr, arr2, 1)
 array([[ 91,  80, 100,  84,  88],
        [ 99, 106, 126,  92,  77],
        [ 69,  98,  91,  93, 117],
        [ 80,  79,  87,  93,  61],
        [ 44,  72,  72,  63,  74]])
 >>> strideConv(arr, arr2, 2)
 array([[ 91, 100,  88],
        [ 69,  91, 117],
        [ 44,  72,  74]])

据我所知,在numpy或scipy中没有支持stride和padding的卷积过滤器的直接实现所以我认为最好使用DL包如torch或tensorflow,然后投射最终结果为numpy。火炬实现可能是:

import torch
import torch.nn.functional as F

arr = torch.tensor(np.expand_dims(arr, axis=(0,1))
arr2 = torch.tensor(np.expand_dims(arr2, axis=(0,1))
output = F.conv2d(arr, arr2, stride=2, padding=0)
output = output.numpy().squeeze()

output>
array([[ 91, 100,  88],
       [ 69,  91, 117],
       [ 44,  72,  74]])