有什么方法可以优化numpy中的循环吗？

Question

我在这里有一些代码（用于梯度计算）-注释了示例值：

dE_dx_strided = np.einsum('wxyd,ijkd->wxyijk', dE_dy, f)
# dE_dx_strided.shape = (64, 25, 25, 4, 4, 3)

imax, jmax, di, dj = dE_dx_strided.shape[1:5]
# imax, jmax, di, dj = (25, 25, 4, 4)

dE_dx = np.zeros_like(x)
# dE_dx.shape = (64, 28, 28, 3)

for i in range(imax):
    for j in range(jmax):
        dE_dx[:, i:i+di, j:j+dj, :] += dE_dx_strided[:, i, j, ...]

其中dE_dx是感兴趣的对象，dE_dx_strided是6张量，有效地被累加在“逐段”上，它看起来像是沿轴{{1}的卷积运算}和1：

2

我最初的想法是使# Verbose convolution operation (not my actual implementation) for i in range(imax): for j in range(jmax): # Vaguely similar, but with filter multiplication, and = instead of += y[i, j] = x[i:i+di, j:dj] * f[di, dj] 的所有要添加到单个dE_dx_strided的元素都沿一个轴放置，然后求和。但我无法使它正常工作。

现在我知道dE_dx[:, i:i+di, j:j+dj, :]循环并不是天生就很慢，但是是否有一种麻木式的方法可以进一步优化此过程，例如通过重塑，求和，跨步等？