求大型向量矩阵点积的最快方法

Question

我正在寻找解决以下问题的最有效方法的建议：

我有两个数组 A 和 B。它们的形状都是 NxNx3。它们表示两个 2D 位置矩阵，其中每个位置都是 x、y 和 z 坐标的向量。

我想创建一个名为 C 的新数组，形状为 NxN，其中 C[i, j] 是向量 A[i, j] 和 B[i, j] 的点积。

以下是我到目前为止提出的解决方案。第一个使用 numpy 的 einsum 函数 (which is beautifully described here)。第二个使用 numpy 的广播规则及其求和函数。

>>> import numpy as np
>>> A = np.random.randint(0, 10, (100, 100, 3))
>>> B = np.random.randint(0, 10, (100, 100, 3))
>>> C = np.einsum("ijk,ijk->ij", A, B)
>>> D = np.sum(A * B, axis=2)
>>> np.allclose(C, D)
True

有没有更快的方法？我听说 numpy 的 tensordot 函数可以非常快地运行，但我一直很难理解它。使用 numpy 的点或内部函数怎么样？

对于某些上下文，A 和 B 数组通常包含 100 到 1000 个元素。

非常感谢任何指导！

Answer 1

稍加整形，我们就可以使用 matmul。这个想法是将前两个维度视为 'batch' 维度，并将最后一个维度 dot 视为：

In [278]: E = A[...,None,:]@B[...,:,None]                                       
In [279]: E.shape                                                               
Out[279]: (100, 100, 1, 1)
In [280]: E = np.squeeze(A[...,None,:]@B[...,:,None])                           
In [281]: np.allclose(C,E)                                                      
Out[281]: True
In [282]: timeit E = np.squeeze(A[...,None,:]@B[...,:,None])                    
130 µs ± 2.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [283]: timeit C = np.einsum("ijk,ijk->ij", A, B)                             
90.2 µs ± 1.53 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

比较时间可能有点棘手。在当前版本中，einsum 可以根据维度采取不同的路线。在某些情况下，它似乎将任务委托给 matmul（或至少是相同的底层类 BLAS 代码）。虽然 einsum 在此测试中更快是件好事，但我不会一概而论。

tensordot 只是重塑（如果需要的话转置）数组，以便它可以应用普通的 2d np.dot。实际上它在这里不起作用，因为您将前 2 个轴视为 'batch'，而它对它们执行 outer product。

求大型向量矩阵点积的最快方法

Fastest Way to Find the Dot Product of a Large Matrix of Vectors

python

numpy

tensordot