批量卷积2d在numpy没有scipy?

时间:2017-08-07 05:35:55

标签: python numpy numpy-einsum

我有一批存储在数组b中的m x n x张图片,以及一张我想要的大小为f的卷积过滤器p x q应用于批处理中的每个图像(然后使用求和池并存储在数组y中),即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))为真。

改编this answer,我可以写下以下内容:

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))

但我认为有一种方法可以只使用一个einsum,这对我很有用,因为b通常在100到1000之间。

如何调整我的方法只使用一个einsum?另外,就我的目的而言,除了scipy之外,我无法引入numpy或任何其他依赖项。

1 个答案:

答案 0 :(得分:3)

只需要将shape设为5d,然后让stridesshape匹配。

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)

如果M变得非常大,现在b可能变得非常大,但它可以解决您的玩具问题。