numpy - 使用 Numpy 沿轴相乘矩阵(点积)(无循环)

Question

我正在使用 Numpy 解决图像处理问题，我试图避免循环并执行以下操作:

我有一个由 NxNxKxK 点组成的矩阵 M(它是由 KxK 矩阵组成的 NxN 矩阵)，对于每一行，我希望将该行中的所有 N 个矩阵 (KxK) 相乘(点积)。因此，如果我对整个 M(所有行)执行此操作，我会得到矩阵 (KxK) 的向量 V (Nx1)，其中 V[i] 保存 M[i,0]xM[i,1 的点积]x...xM[i,N-1]。

我使用循环实现了这个问题的解决方案，但我无法找到一种无需循环即可实现此目的的方法。

实现(带循环):

    a = np.array([[1,1,1], [1,1,1], [1,1,1]])
    mat = np.array([[a,a,a,a], [a*2,a*2,a*2,a*2], [a*3,a*3,a*3,a*3],
                    [a*4,a*4,a*4,a*4]])  # the original matrix
    N, N, k, k = mat.shape
    result = np.ones((N, k, k))  # resulting matrix
    for i in range(N):
        k = functools.reduce(np.dot, mat[i,:])
        result[i,:] = k
    print(result)

Answer 1

以下使用 reduce 但不使用 N 上的循环:

mat = mat.swapaxes(0, 1)
result = functools.reduce(lambda a, b: np.einsum('ijk,ikl->ijl', a, b), mat[:])

einsum 表示法 'jk,kl->jl' 表示矩阵乘法，索引 i 表示应该对每个矩阵乘法进行第一个索引的值。 mat[0] 或 mat[1] 的第一个索引实际上是 mat 的第二个索引(列索引)，因此可以这样写:乘法发生在 mat 的每一列中。您希望在每一行中完成此操作，因此使用了 swapaxes。

这比 for 循环版本更快或更慢取决于 N 和 k 的相对大小。 np.dot 方法经过高度优化，但如果 N 上的循环很长，einsum 可能会获胜。一些 %timeit 结果:

• 当 N=100、k=2 时，for 循环版本需要 7.5 ms，einsum 版本需要 4.31 ms。
• 当 N=100、k=20 时，for 循环版本需要 27.3 毫秒，einsum 版本需要 153 毫秒。

因此，在特定情况下会有适度的 yield ，而在大多数其他情况下会出现重大损失。但是您并没有要求一种“高效”的解决方案，您要求的是一种“无循环”的解决方案，所以这就是(“无循环”!=“更快”)。正如 Divakar 在评论中建议的那样，您最好保持代码不变。