为什么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->i'）。它还可能有 2 个和 3 个变量的特殊情况（而不是更多）。但它不会尝试调用外部库。

我制定了一个纯Python的类似工作，但更关注解析而不是计算特殊情况。

tensordot重塑和交换，因此它可以调用dot实际计算。