Question

我有一个512x512的图像阵列，我想对8x8的块执行操作。此刻我有这样的东西：

const IVector2<int> IElevatorHandler::floors(-1, 4);

其中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是从原始数组派生的8x8块。我想通过向量化操作来加快此代码的速度。我已经像这样修改了输入内容：

如何仅在轴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]))，但是我没有遇到一个简单的例子来说明我要实现的目标。

Answer 1

如果您重新订购它们，它们将起作用。如果input和some_other_array的形状均为（64,8,64,8），则：

input = input.transpose(0,2,1,3)
some_other_array = some_other_array.transpose(0,2,1,3)

这会将它们重新排序为64、64、8、8。此时，您可以计算矩阵乘法。请注意，您需要使用matmul来计算块积，而不是点，这会尝试将整个矩阵相乘。

output = input @ some_other_array
output = output.transpose(0,2,1,3)
output = output.reshape(512,512)

Answer 2

您怀疑，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。

沿特定轴的块状点积

2 个答案: