2d batch convolution in numpy without scipy?

Question

I have a batch of images b mxn stored in an array x , and a convolutional filter f size pxq , which I would like to apply to each image (then use summation and store in array y ) in a package, i.e. all(np.allclose(y[i][j][k], (x[i, j:j+p, k:k+q] * f).sum()) for i in range(b) for j in range(m-p+1) for k in range(n-q+1)) is true.

Adapting this answer , I could write the following:

 b, m, n, p, q = 6, 5, 4, 3, 2 x = np.arange(b*m*n).reshape((b, m, n)) f = np.arange(p*q).reshape((p, q)) y = [] for i in range(b): shape = f.shape + tuple(np.subtract(x[i].shape, f.shape) + 1) strides = x[i].strides * 2 M = np.lib.stride_tricks.as_strided(x[i], shape=shape, strides=strides) y.append(np.einsum('ij,ijkl->kl', f, M)) assert all(np.allclose(y[i][j][k], (x[i, j:j+p, k:k+q] * f).sum()) for i in range(b) for j in range(m-p+1) for k in range(n-q+1)) 

but I think there is a way to do this with only one einsum , which would be useful for me, because b usually between 100 and 1000.

How to adapt my approach to using only one einsum ? Also, for my purposes, I cannot add scipy or any other dependencies other than numpy .

Answer 1

You just need to get a shape like 5d and get strides according to shape .

 shape = f.shape + (x.shape[0],) + tuple(np.subtract(x.shape[1:], f.shape) + 1) strides = (x.strides * 2)[1:] M = np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides) y = np.einsum('pq,pqbmn->bmn', f, M) 

now M can get really big if b gets really big, but it works on your toy problem.