点积两个 4D Numpy 数组

Question

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

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

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 在 cmets 中建议，您可以使用下面提到的代码的较短版本。此代码将 ij 与 A 元素相乘，然后将其与 i（又名行）相加并返回（相当于得到点积的对角线）：

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() 仅返回其输入数组的视图，不进行任何昂贵的数据复制。
• 这里不需要初始化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 答案。）