如何从 np.tensordot 到 np.einsum

【问题标题】：How to go from np.tensordot to np.einsum如何从 np.tensordot 到 np.einsum
【发布时间】：2017-03-23 14:20:10
【问题描述】：

我收到的代码中的内容类似于：

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

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

我应该改变，因为有人告诉我使用np.einsum 会更快，但我想我不完全理解tensordot 输出的内容。

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

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

【问题讨论】：

标签： python numpy-einsum

【解决方案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

【讨论】：