python - 沿特定轴的 Numpy 点积

Question

我有一个 512x512 的图像数组，我想对 8x8 的 block 执行操作。目前我有这样的事情:

output = np.zeros(512, 512)

for i in range(0, 512, 8):
    for j in rangerange(0, 512, 8):
        a = input[i:i+8, j:j+8]
        b = some_other_array[i:i+8, j:j+8]
        output[i:i+8, j:j+8] = np.dot(a, b)

其中 a 和 b 是从原始数组派生的 8x8 block 。我想通过使用向量化操作来加速这段代码。我已经像这样 reshape 了我的输入:

input = input.reshape(64, 8, 64, 8)
some_other_array = some_other_array.reshape(64, 8, 64, 8)

我如何仅在轴 1 和 3 上执行点积以输出形状数组 (64, 8, 64, 8)?

我试过 np.tensordot(input, some_other_array, axes=([0, 1], [2, 3])) 给出了正确的输出形状，但值不匹配上面循环的输出。我也查看了 np.einsum，但我还没有遇到一个简单的例子来说明我想要实现的目标。

Answer 1

如您所料，np.einsum 可以解决这个问题。如果 input 和 some_other_array 的形状是 (64, 8, 64, 8)，那么如果你写

output = np.einsum('ijkl,ilkm->ijkm', input, some_other_array)  

然后 output 也将具有形状 (64, 8, 64, 8)，其中矩阵乘法(即 np.dot)具有仅在轴 1 和 3 上完成。

np.einsum 的字符串参数看起来很复杂，但实际上它是两件事的结合。首先，矩阵乘法由 jl,lm->jm 给出(参见例如 this answer on einsum )。其次，我们不想对轴 0 和 2 做任何事情，所以对于它们我只写了 ik,ik->ik。将两者结合得到 ijkl,ilkm->ijkm。