时间戳

首页 > 解决方案 > 执行大点/张量点积的最有效方法，同时只保留对角线条目

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

问题描述

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

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

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

第一种方法：

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)矩阵，并且也花费了太多时间，我想知道是否有任何有效的方法可以做到这一点？

标签： pythonnumpylinear-algebratensordot

解决方案

定义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。

推荐阅读