Strided convolution of 2D in numpy

Question

I tried to implement strided convolution of a 2D array using for loop i.e

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)

This results in 3*3 array:

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

Is there a numpy function or scipy function to do the same? My approach is not that good. How can I vectorize this?

Answer 1

Ignoring the padding argument and trailing windows that won't have enough lengths for convolution against the second array, here's one way with 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)))

Alternatively, we can use the scikit-image built-in view_as_windows to get those windows elegantly, like so -

from skimage.util.shape import view_as_windows

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