点积两个 4D Numpy 数组

Question

我有一个形状为 (15, 2, 320, 320) 的 4D Numpy 数组。换句话说，[320 x 320] 矩阵的每个元素都是大小为 [15 x 2] 的矩阵。现在，我想计算 [320x320] 矩阵的每个元素的点积，然后提取对角线数组。

目前，我正在使用 2 个“for”循环，并且代码运行良好（请参阅代码片段）。但是，计算速度太慢（当我处理大数据时）。任何人都可以告诉我如何在没有循环的情况下矢量化计算。

A = np.random.rand(15, 2, 320, 320)
B = np.zeros((2, 320, 320))  # store the final results
 
for row in range (320):
        for col in range (320):
            C = np.dot(A[:, :, row, col].T, A[:, :, row, col]) # the shape of C is [2 x 2]
            C = np.diag(C)
            B[:, row, col] = C

如有任何帮助或建议，我们将不胜感激！

Answer 1

编辑：@hpaulj 在评论中建议，您可以使用下面提到的代码的较短版本。此代码将 A 的 ij 逐元素相乘，然后对 i(a.k.a.rows) 和 returns 求和（相当于获取点的对角线产品):

np.einsum('ijkl,ijkl->jkl',A,A)

您可以使用一行代码对其进行矢量化，如下所示。当然，您可以使用转置和点积的组合来替代 np.einsum，但我认为 np.einsum 中的索引使代码更清晰，更易读。 np.diagonal 采用前两个轴的对角线作为最后一个维度，因此，您需要在末尾进行转置以将该维度带到您想要的左侧：

np.diagonal(np.einsum('ijkl,imkl->jmkl',A,A)).transpose(2,0,1)

测试输出是否与您的代码输出匹配 B:

np.allclose(np.diagonal(np.einsum('ijkl,imkl->jmkl',A,A)).transpose(2,0,1), B)
#True

Answer 2

这应该有效：

A_T_1 = np.transpose(A, (2,3,0,1))           # Shape (320,320,2,15)
A_T_2 = np.transpose(A, (2,3,1,0))           # Shape (320,320,15,2)
C_all = A_T_2 @ A_T_1                        # Shape (320,320,2,2)
D_all = np.diagonal(C_all, axis1=2, axis2=3) # Shape (320,320,2)

B = np.transpose(D_all, (2,0,1))             # Shape (2,320,320)

注：

transpose() 仅 returns 其输入数组的视图，并且不进行任何昂贵的数据复制。
这里不需要初始化B = np.zeros((2, 320, 320))

Answer 3

有时张量代数的计算方式有点有趣。只是 (A*A).sum(0):

A = np.random.rand(15, 2, 320, 320)
B1 = np.zeros((2, 320, 320))  # store the final results
 
for row in range (320):
        for col in range (320):
            C = np.dot(A[:, :, row, col].T, A[:, :, row, col]) # the shape of C is [2 x 2]
            C = np.diag(C)
            B1[:, row, col] = C

B2 = (A*A).sum(0)

np.allclose(B1, B2)
True

（感谢@Ehsan 和@hpaulj，他们提出了 einsum 的答案，提出了这个公式。）

点积两个 4D Numpy 数组

Dot product two 4D Numpy array

python

arrays

numpy

matrix-multiplication