让我们从三个
dtype=np.double
数组开始。计时是在 intel CPU 上使用 numpy 1.7.1 进行的,该版本使用 icc
编译并链接到 intel 的 mkl
。还使用使用 gcc
编译而没有 mkl
的 numpy 1.6.1 AMD cpu 来验证时序。请注意,时间比例与系统大小几乎呈线性关系,并不是由于 numpy 函数 if
语句中产生的小开销造成的,这些差异将以微秒而不是毫秒显示:
arr_1D=np.arange(500,dtype=np.double)
large_arr_1D=np.arange(100000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)
首先让我们看一下
np.sum
函数:
np.all(np.sum(arr_3D)==np.einsum('ijk->',arr_3D))
True
%timeit np.sum(arr_3D)
10 loops, best of 3: 142 ms per loop
%timeit np.einsum('ijk->', arr_3D)
10 loops, best of 3: 70.2 ms per loop
权力:
np.allclose(arr_3D*arr_3D*arr_3D,np.einsum('ijk,ijk,ijk->ijk',arr_3D,arr_3D,arr_3D))
True
%timeit arr_3D*arr_3D*arr_3D
1 loops, best of 3: 1.32 s per loop
%timeit np.einsum('ijk,ijk,ijk->ijk', arr_3D, arr_3D, arr_3D)
1 loops, best of 3: 694 ms per loop
外在产品:
np.all(np.outer(arr_1D,arr_1D)==np.einsum('i,k->ik',arr_1D,arr_1D))
True
%timeit np.outer(arr_1D, arr_1D)
1000 loops, best of 3: 411 us per loop
%timeit np.einsum('i,k->ik', arr_1D, arr_1D)
1000 loops, best of 3: 245 us per loop
使用
np.einsum
,上述所有速度都快了一倍。这些应该是同类比较,因为一切都是特定的dtype=np.double
。我希望这样的操作能够加快速度:
np.allclose(np.sum(arr_2D*arr_3D),np.einsum('ij,oij->',arr_2D,arr_3D))
True
%timeit np.sum(arr_2D*arr_3D)
1 loops, best of 3: 813 ms per loop
%timeit np.einsum('ij,oij->', arr_2D, arr_3D)
10 loops, best of 3: 85.1 ms per loop
无论 np.inner
选择如何,Einsum 似乎至少比
np.outer
、np.kron
、np.sum
和 axes
快两倍。主要的例外是 np.dot
,因为它从 BLAS 库调用 DGEMM。那么为什么 np.einsum
比其他等效的 numpy 函数更快呢?
DGEMM 案例的完整性:
np.allclose(np.dot(arr_2D,arr_2D),np.einsum('ij,jk',arr_2D,arr_2D))
True
%timeit np.einsum('ij,jk',arr_2D,arr_2D)
10 loops, best of 3: 56.1 ms per loop
%timeit np.dot(arr_2D,arr_2D)
100 loops, best of 3: 5.17 ms per loop
主要理论来自 @sebergs 评论,
np.einsum
可以利用SSE2,但 numpy 的 ufuncs 直到 numpy 1.8 才会使用(请参阅更改日志)。我相信这是正确的答案,但不能能够确认它。通过更改输入数组的数据类型并观察速度差异以及并非每个人都观察到相同的时序趋势这一事实,可以找到一些有限的证据。
首先,过去在 numpy 列表上对此有很多讨论。例如,参见:
其中一些归结为以下事实:
einsum
是新的,并且可能试图更好地处理缓存对齐和其他内存访问问题,而许多旧的 numpy 函数专注于易于移植的实现而不是高度优化的实现。不过,我只是猜测。
但是,您正在做的一些事情并不完全是“同类”比较。
除了@Jamie已经说过的之外,
sum
对数组使用了更合适的累加器
例如,
sum
更仔细地检查输入的类型并使用适当的累加器。例如,请考虑以下情况:
In [1]: x = 255 * np.ones(100, dtype=np.uint8)
In [2]: x
Out[2]:
array([255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255, 255, 255, 255, 255], dtype=uint8)
请注意,
sum
是正确的:
In [3]: x.sum()
Out[3]: 25500
虽然
einsum
会给出错误的结果:
In [4]: np.einsum('i->', x)
Out[4]: 156
但是如果我们使用限制较少的
dtype
,我们仍然会得到您期望的结果:
In [5]: y = 255 * np.ones(100)
In [6]: np.einsum('i->', y)
Out[6]: 25500.0
现在 numpy 1.8 已经发布,根据文档,所有 ufunc 都应该使用 SSE2,我想仔细检查 Seberg 关于 SSE2 的评论是否有效。
为了执行测试,创建了新的 python 2.7 安装 - numpy 1.7 和 1.8 是在运行 Ubuntu 的 AMD opteron 核心上使用标准选项通过
icc
进行编译的。
这是1.8升级前后的测试运行:
import numpy as np
import timeit
arr_1D=np.arange(5000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)
print 'Summation test:'
print timeit.timeit('np.sum(arr_3D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print timeit.timeit('np.einsum("ijk->", arr_3D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print '----------------------\n'
print 'Power test:'
print timeit.timeit('arr_3D*arr_3D*arr_3D',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print timeit.timeit('np.einsum("ijk,ijk,ijk->ijk", arr_3D, arr_3D, arr_3D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print '----------------------\n'
print 'Outer test:'
print timeit.timeit('np.outer(arr_1D, arr_1D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print timeit.timeit('np.einsum("i,k->ik", arr_1D, arr_1D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print '----------------------\n'
print 'Einsum test:'
print timeit.timeit('np.sum(arr_2D*arr_3D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print timeit.timeit('np.einsum("ij,oij->", arr_2D, arr_3D)',
'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
number=5)/5
print '----------------------\n'
Numpy 1.7.1:
Summation test:
0.172988510132
0.0934836149216
----------------------
Power test:
1.93524689674
0.839519000053
----------------------
Outer test:
0.130380821228
0.121401786804
----------------------
Einsum test:
0.979052495956
0.126066613197
Numpy 1.8:
Summation test:
0.116551589966
0.0920487880707
----------------------
Power test:
1.23683619499
0.815982818604
----------------------
Outer test:
0.131808176041
0.127472200394
----------------------
Einsum test:
0.781750011444
0.129271841049
我认为这是相当结论性的,SSE 在时序差异中发挥了很大的作用,应该注意的是,重复这些测试的时序仅约 0.003 秒。剩余的差异应在该问题的其他答案中涵盖。
我认为这些时间安排可以解释正在发生的事情:
a = np.arange(1000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 3.32 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 6.84 us per loop
a = np.arange(10000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 12.6 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 16.5 us per loop
a = np.arange(100000, dtype=np.double)
%timeit np.einsum('i->', a)
10000 loops, best of 3: 103 us per loop
%timeit np.sum(a)
10000 loops, best of 3: 109 us per loop
因此,当通过
np.sum
调用 np.einsum
时,基本上有一个几乎恒定的 3us 开销,因此它们基本上运行得很快,但需要更长的时间才能开始。为什么会这样呢?我的钱花在以下方面:
a = np.arange(1000, dtype=object)
%timeit np.einsum('i->', a)
Traceback (most recent call last):
...
TypeError: invalid data type for einsum
%timeit np.sum(a)
10000 loops, best of 3: 20.3 us per loop
不确定到底发生了什么,但似乎
np.einsum
正在跳过一些检查来提取类型特定函数来执行乘法和加法,并且仅针对标准 C 类型直接使用 *
和 +
。
多维情况没有不同:
n = 10; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
100000 loops, best of 3: 3.79 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 7.33 us per loop
n = 100; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
1000 loops, best of 3: 1.2 ms per loop
%timeit np.sum(a)
1000 loops, best of 3: 1.23 ms per loop
因此,开销基本上是恒定的,一旦他们开始做,就不会跑得更快。
numpy 1.21.2 的更新:Numpy 的本机函数在几乎所有情况下都比 einsum 更快。只有 einsum 的外部变体和 sum23 测试速度比非 einsum 版本更快。
如果你可以使用 numpy 的本机函数,就这样做。
(使用我的项目perfplot创建的图像。)
重现绘图的代码:
import numpy
import perfplot
def setup1(n):
return numpy.arange(n, dtype=numpy.double)
def setup2(n):
return numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n)
def setup3(n):
return numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n)
def setup23(n):
return (
numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n),
numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n),
)
def numpy_sum(a):
return numpy.sum(a)
def einsum_sum(a):
return numpy.einsum("ijk->", a)
perfplot.save(
"sum.png",
setup=setup3,
kernels=[numpy_sum, einsum_sum],
n_range=[2 ** k for k in range(10)],
)
def numpy_power(a):
return a * a * a
def einsum_power(a):
return numpy.einsum("ijk,ijk,ijk->ijk", a, a, a)
perfplot.save(
"power.png",
setup=setup3,
kernels=[numpy_power, einsum_power],
n_range=[2 ** k for k in range(9)],
)
def numpy_outer(a):
return numpy.outer(a, a)
def einsum_outer(a):
return numpy.einsum("i,k->ik", a, a)
perfplot.save(
"outer.png",
setup=setup1,
kernels=[numpy_outer, einsum_outer],
n_range=[2 ** k for k in range(13)],
)
def dgemm_numpy(a):
return numpy.dot(a, a)
def dgemm_einsum(a):
return numpy.einsum("ij,jk", a, a)
def dgemm_einsum_optimize(a):
return numpy.einsum("ij,jk", a, a, optimize=True)
perfplot.save(
"dgemm.png",
setup=setup2,
kernels=[dgemm_numpy, dgemm_einsum],
n_range=[2 ** k for k in range(13)],
)
def dot_numpy(a):
return numpy.dot(a, a)
def dot_einsum(a):
return numpy.einsum("i,i->", a, a)
perfplot.save(
"dot.png",
setup=setup1,
kernels=[dot_numpy, dot_einsum],
n_range=[2 ** k for k in range(20)],
)
def sum23_numpy(data):
a, b = data
return numpy.sum(a * b)
def sum23_einsum(data):
a, b = data
return numpy.einsum("ij,oij->", a, b)
perfplot.save(
"sum23.png",
setup=setup23,
kernels=[sum23_numpy, sum23_einsum],
n_range=[2 ** k for k in range(10)],
)