如何从np.tensordot转到np.einsum

Question

我得到的代码是这样的：

C = np.tensordot(B, A, axes = (0,0))

A是(20L, 50L)，B是(20L, 20L)

自从有人告诉我使用np.einsum会更快之后，我应该改变一下，但我想我不完全理解tensordot输出的是什么。

现在C是一个(20L, 50L)，我不明白为什么？

当然，我已经阅读了文档页面，但仍然没有掌握其中的信息。考虑到我刚刚开始使用Python。

Answer 1

其中A具有形状(20,50)和B具有形状(20,20)：

C = np.tensordot(B, A, axes = (0,0))

等同于：

C = np.zeros((20,50))
for i in range(20):
    for j in range(50):
        for k in range(20):
            C[i,j] += A[k,j] * B[k,i]

相当于：

C = np.einsum('kj,ki->ij',A,B)

einsum看起来确实比tensordot快一点：

%timeit np.einsum('kj,ki->ij',A,B)
10000 loops, best of 3: 20.1 µs per loop

%timeit np.tensordot(B,A,axes=(0,0))
10000 loops, best of 3: 33.6 µs per loop