python - 为什么 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实际计算。