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

Question

我有使用OpenBlas编译的numpy，我想知道为什么einsum比dot慢得多（我理解在3个索引的情况下，但我不明白为什么它在两个索引情况下的性能也不太好）？这是一个例子：

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使用OpenBlas和numpy.dot这样的并行化？ 为什么np.einsum不只是在注意到点积时调用np.dot？

Answer 1

einsum解析索引字符串，然后构造一个nditer对象，并使用它来执行产品和迭代。它具有特殊情况，其中索引仅执行轴交换，并且求和（'ii-＆gt; i'）。它也可能有2和3个变量的特殊情况（而不是更多）。但是它没有尝试调用外部库。

我制作了一个纯粹的python work-a-like，但是比计算特殊情况更侧重于解析。

tensordot重新整形和交换，然后可以调用dot进行实际计算。