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

Question

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

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

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

这是我迄今为止提出的解决方案。第一个使用 numpy 的 einsum 函数 (which is beautifully described here)。第二个使用 numpy 的广播规则及其 sum 函数。

>>> 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。我们的想法是将前 2 个维度视为“批量”维度，最后一个维度为 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 个轴视为一个“批次”，就像它对它们执行 outer product 一样。