python - 执行大点/张量点积同时仅保留对角线条目的最有效方法

Question

这个问题已经有答案了:

Matrix multiplication for multidimensional matrix (/array) - how to avoid loop? (3 个回答)

已关闭 3 年前。

我正在尝试找出一种使用 numpy 以尽可能最省时的方式执行以下代数的方法:

给定 3D 矩阵/张量，A ，形状 (n, m, p)和 2D 矩阵/张量 B ，形状 (n, p) ，计算C_ij = sum_over_k (A_ijk * B_ik) ，其中生成的矩阵 C维度为 (n, m)。

我尝试了两种方法来做到这一点。一种是循环遍历第一个维度，并每次计算常规点积。另一种方法是使用np.tensordot(A, B.T)计算形状 (n, m, n) 的结果，然后沿第一维和第三维取对角线元素。两种方法如下所示。

第一种方法:

C = np.zeros((n,m))

for i in range(n):

  C[i] = np.dot(A[i], B[i])

第二种方法:

C = np.diagonal(np.tensordot(A, B.T, axes = 1), axis1=0, axis2=2).T

但是，由于n是一个非常大的数字，第一种方法中n的循环会花费大量时间。第二种方法计算了太多不必要的条目来获得那么大的 (n, m, n)矩阵，而且还花费太多时间，我想知道是否有任何有效的方法来做到这一点？

Answer 1

定义2个数组:

In [168]: A = np.arange(2*3*4).reshape(2,3,4); B = np.arange(2*4).reshape(2,4)                               

您的迭代方法:

In [169]: [np.dot(a,b) for a,b in zip(A,B)]                                                                  
Out[169]: [array([14, 38, 62]), array([302, 390, 478])]

einsum 实际上是从您的 C_ij = sum_over_k (A_ijk * B_ik) 中自行编写的:

In [170]: np.einsum('ijk,ik->ij', A, B)                                                                      
Out[170]: 
array([[ 14,  38,  62],
       [302, 390, 478]])

添加了

@、matmul来执行批量dot产品；这里的i维度是批处理维度。由于它使用 A 的最后一个和 B 的倒数第二个进行 dot 求和，因此我们必须暂时扩展 B 到 (2,4,1):

In [171]: A@B[...,None]                                                                                      
Out[171]: 
array([[[ 14],
        [ 38],
        [ 62]],

       [[302],
        [390],
        [478]]])
In [172]: (A@B[...,None])[...,0]                                                                             
Out[172]: 
array([[ 14,  38,  62],
       [302, 390, 478]])

通常 matmul 是最快的，因为它将任务传递给类似 BLAS 的代码。