为什么numpy.dot比numpy.einsum快得多？

Question

我已经使用OpenBlas编译了numpy，我想知道为什么einsum比dot慢得多（我理解在三个索引的情况下，但我不明白为什么在两个索引的情况下性能也较差）？这里有一个例子：

import numpy as np
A = np.random.random([1000,1000])
B = np.random.random([1000,1000])

%timeit np.dot(A,B)

Out: 10 loops, best of 3: 26.3 ms per loop

%timeit np.einsum("ij,jk",A,B)

Out: 5 loops, best of 3: 477 ms per loop

有没有办法让einsum像numpy.dot一样使用OpenBlas和并行化呢？为什么np.einsum如果注意到点积不直接调用np.dot呢？

- varantir

---

相关链接：https://dev59.com/y2Ij5IYBdhLWcg3wilkI - Alex Riley

这是正确的，但我的问题更与两个索引有关。给定帖子中的答案让我觉得三个索引存在问题。我对这个特定例子中速度差异如此惊讶。 - varantir

einsum被设计成通用的，并没有特殊处理只有两个索引的点积使用blas。这就是实质问题。如果你知道只有两个索引，那就直接使用dot。 - BeRecursive

Answer 1

einsum解析索引字符串，然后构造一个nditer对象，并使用它执行积和求和迭代。它有一些特殊情况，其中索引只执行轴交换和求和（'ii->i'）。它还可能有针对2个和3个变量的特殊情况（而不是更多）。但它不会尝试调用外部库。

我设计了一个纯Python的类似实现，但更加关注解析而不是计算特殊情况。 tensordot重新整形和交换，因此它可以调用dot进行实际计算。

- hpaulj