np.einsum 中的速度差异答案

【问题标题】：Speed difference in np.einsumnp.einsum 中的速度差异
【发布时间】：2017-08-16 14:37:25
【问题描述】：

我注意到np.einsum 在减少一维时速度更快

import numpy as np
a = np.random.random((100,100,100))
b = np.random.random((100,100,100))

%timeit np.einsum('ijk,ijk->ijk',a,b)
# 100 loops, best of 3: 3.83 ms per loop
%timeit np.einsum('ijk,ijk->ij',a,b)
# 1000 loops, best of 3: 937 µs per loop
%timeit np.einsum('ijk,ijk->i',a,b)
# 1000 loops, best of 3: 921 µs per loop
%timeit np.einsum('ijk,ijk->',a,b)
# 1000 loops, best of 3: 928 µs per loop

这对我来说似乎很奇怪，因为我希望它首先生成新数组然后对其求和，这显然不会发生。那里发生了什么？为什么下降一个维度会变快，而下降另一个维度不会变快？

旁注：我首先认为它与创建一个大数组有关，当它有很多维度时，我认为不是这样：

 %timeit np.ones(a.shape)
 # 1000 loops, best of 3: 1.79 ms per loop
 %timeit np.empty(a.shape)
 # 100000 loops, best of 3: 3.05 µs per loop

因为创建新数组要快得多。

【问题讨论】：

标签： python numpy numpy-einsum

【解决方案1】：

einsum 在编译代码中实现，numpy/numpy/core/src/multiarray/einsum.c.src。

核心操作是使用nditer 的c 版本迭代所有维度（例如在您的情况下为100*100*100 次），应用ijk 字符串定义的sum-of-products 计算。

但它会进行各种优化，包括在不需要乘法时生成视图。因此，您需要仔细研究以了解您的案例有何不同。

时间划分是在产生没有求和的 3d 输出和在一个或多个轴上求和的输出之间。

【讨论】：