执行大点/张量点积的最有效方法，同时只保留对角线条目[重复]答案

【问题标题】：Most efficient way to perform large dot/tensor dot products while only keeping diagonal entries [duplicate]执行大点/张量点积的最有效方法，同时只保留对角线条目[重复]
【发布时间】：2020-01-09 19:41:57
【问题描述】：

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

给定一个形状为(n, m, p) 的3D 矩阵/张量A 和一个形状为(n, p) 的2D 矩阵/张量B，计算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)matrix，而且也花费了太多的时间，我想知道是否有任何有效的方法可以做到这一点？

【问题讨论】：

标签： python numpy linear-algebra tensordot

【解决方案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。

【讨论】：

【解决方案2】：

这是我的实现：

B = np.expand_dims(B, axis=1)
E = A * B
E = np.sum(E, axis=-1)

检查：

import numpy as np
n, m, p = 2, 2, 2
np.random.seed(0)
A = np.random.randint(1, 10, (n, m, p))
B = np.random.randint(1, 10, (n, p))

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

# from here is my implementation
B = np.expand_dims(B, axis=1)
E = A * B
E = np.sum(E, axis=-1)

print(np.array_equal(C, E))

True

使用np.expand_dims() 添加新维度。并使用广播乘法。最后，沿第三维求和。

感谢user3483203的验证码

【讨论】：