将双重（三重）求和写成内积？

Question

由于我的np.dot受到OpenBlas和Openmpi的加速，我想知道是否有可能编写双重求和。

for i in range(N):
     for j in range(N):
         B[k,l] += A[i,j,k,l] * X[i,j]

作为内积。现在我正在使用。

B = np.einsum("ijkl,ij->kl",A,X)

但不幸的是它运行速度非常慢，并且只使用一个处理器。有什么好的想法吗？

编辑：我已经用一个简单的例子对给出的答案进行了基准测试，似乎它们都处于同一数量级：

A = np.random.random([200,200,100,100])
X = np.random.random([200,200])
def B1():
    return es("ijkl,ij->kl",A,X) 
def B2():
    return np.tensordot(A, X, [[0,1], [0, 1]])
def B3():
    shp = A.shape
    return np.dot(X.ravel(),A.reshape(shp[0]*shp[1],1)).reshape(shp[2],shp[3])

%timeit B1()
%timeit B2()
%timeit B3()

1 loops, best of 3: 300 ms per loop
10 loops, best of 3: 149 ms per loop
10 loops, best of 3: 150 ms per loop

从这些结果来看，我会选择np.einsum，因为它的语法仍然是最易读的，而其他两个选项的改进只有2倍的因素。我猜下一步是将代码外部化为C或Fortran。

- varantir

---

它有多慢？你有在上面使用任何基准测试吗？如果您能在此处发布完整代码，那将很有意义。for循环的第一部分非常容易理解，但最后一行并没有帮助我们理解完整的流程。请发布完整的代码。 - Pralhad Narsinh Sonar

亲爱的Pralhad，这是完整的代码，你可以选择任何ndarray A、X。本文的目的不是解释einsum函数，而是找到一个用np.dot表示的表达式，因为np.einsum没有并行化。没有必要给出基准，因为我没有什么可以与np.einsum进行比较的（虽然我可以对for循环进行基准测试，但由于它是本地Python代码，速度慢了几个数量级，因此没有用处）。 - varantir

np.einsum 在 C 中已经进行了优化并使用了 SSE 向量化，因此即使在单线程执行中使用 OpenBlas，您也无法获得数量级的速度提升。使用所提出的方法可以实现 2 倍的加速是合理的。 - rth

虽然EDIT中的运行时间似乎合理，但是np.dot方法中有一个错别字：A.reshape(shp[0]*shp[1],1)，其中应该使用-1。 - Divakar

当对第二个矩阵的所有维度求和（就像在这里所做的那样）时，使用tensordot几乎没有什么好处。然而，当您实际执行张量积时（即当两个矩阵都具有不被求和的维度时），tensordot显示出巨大的优势。另一种说法是，如果您可以将问题“重塑”为矩阵-向量乘积，则einsum的速度与tensordot相同。但是，如果您只能将问题“重塑”为矩阵-矩阵乘积，则tensordot将更快。 - Will Martin

Answer 1

你可以使用np.tensordot()函数：

np.tensordot(A, X, [[0,1], [0, 1]])

这段代码使用了多核心。

---

编辑：有趣的是，当增加输入数组的大小时，np.einsum和np.tensordot的扩展性如何：

In [18]: for n in range(1, 31):
   ....:     A = np.random.rand(n, n+1, n+2, n+3)
   ....:     X = np.random.rand(n, n+1)
   ....:     print(n)
   ....:     %timeit np.einsum('ijkl,ij->kl', A, X)
   ....:     %timeit np.tensordot(A, X, [[0, 1], [0, 1]])
   ....:
1
1000000 loops, best of 3: 1.55 µs per loop
100000 loops, best of 3: 8.36 µs per loop
...
11
100000 loops, best of 3: 15.9 µs per loop
100000 loops, best of 3: 17.2 µs per loop
12
10000 loops, best of 3: 23.6 µs per loop
100000 loops, best of 3: 18.9 µs per loop
...
21
10000 loops, best of 3: 153 µs per loop
10000 loops, best of 3: 44.4 µs per loop

当使用tensordot处理更大的数组时，其优势变得更加明显。

- Saullo G. P. Castro

---

你使用了多少个核心？在numpy下使用什么库？我有8个核心，并使用MKL。我记得通过切换到tensordot获得了更好的改进。 - Will Martin

我发现我的电脑使用了4个内核，并且使用了MKL NumPy...他们在einsum中使用SIMD编程进行了一些改进，使它更具竞争力。 - Saullo G. P. Castro

Answer 2

您可以像这样使用 np.dot -

shp = A.shape
B_dot = np.dot(X.ravel(),A.reshape(shp[0]*shp[1],-1)).reshape(shp[2],shp[3])

- Divakar

Answer 3

我发现在某些操作中，tensordot 的性能远高于 einsum。我正在使用安装了加速器的 Anaconda 中的 numpy，并且安装了英特尔的数学核心库（MKL）。考虑当第二个矩阵有一个额外维度未被求和时会发生什么：

In [39]: A = np.random.random([200, 200, 100, 100])

In [40]: X = np.random.random([200, 200])

In [41]: Y = np.random.random([200, 200, 100])

我正在进行的操作如下：

A X ---> (100, 100)

A Y ---> (100, 100, 100)

在此设置中，A Y 操作基本上需要执行 100 次 A X 操作，并存储每个操作的结果。张量点积在这种情况下的执行方式如下：

In [42]: %timeit tensordot(A, X, [(0,1), (0,1)])
1 loops, best of 3: 477 ms per loop

In [43]: %timeit tensordot(A, Y, [(0,1), (0,1)])
1 loops, best of 3: 1.35 s per loop

停下来思考一下。在第43行，我们刚刚执行了100倍的操作，但只需要3倍的时间。我知道MKL会使用一些花哨的CPU缓存技术来避免从RAM传输数据。我猜它正在重复使用A的块来处理额外的100个Y数组。

使用Einsum可以得到更符合预期的结果，因为我们需要执行100倍的操作：

In [44]: %timeit einsum('ijkl,ij->kl', A, X)
1 loops, best of 3: 962 ms per loop

In [45]: %timeit einsum('ijkl,ijm->klm', A, Y)
1 loops, best of 3: 1min 45s per loop

似乎当其中一个参数数组的所有维度都被求和时，einsum表现得非常出色。 当一些维度没有被求和时（类似于np.outer但是用于多维数组），使用tensordot可以获得巨大的性能提升。

以下是另一个例子：

对于数组操作：

50x1000x1000 X 50x1000x1000 -> 50x50

使用tensordot给我提供了6 GFLOPS，而einsum只有0.2 GFLOPS。

我认为一个重要的观点是现代计算机应该能够在大型数组中达到5-50 GFLOP的范围。如果您计算操作次数并且得到的结果少于这个值，请检查您正在使用的库。

- Will Martin