numpy.sum可能比Python for循环慢

时间:2013-01-28 16:30:04

标签: python performance numpy vectorization

在特定轴上求和数组时,专用数组方法array.sum(ax)实际上可能比for循环慢:

v = np.random.rand(3,1e4)

timeit v.sum(0)                             # vectorized method
1000 loops, best of 3: 183 us per loop

timeit for row in v[1:]: v[0] += row        # python loop
10000 loops, best of 3: 39.3 us per loop

矢量化方法比普通的for循环慢4倍!这里(g)的(wr)是什么,我不能相信numpy中的矢量化方法比for循环更快吗?

1 个答案:

答案 0 :(得分:8)

不,你不能。正如您的有趣示例所指出的那样numpy.sum可能不是最理想的,并且通过显式for循环更好地布局操作可以更有效。

让我再举一个例子:

>>> N, M = 10**4, 10**4
>>> v = np.random.randn(N,M)
>>> r = np.empty(M)
>>> timeit.timeit('v.sum(axis=0, out=r)', 'from __main__ import v,r', number=1)
1.2837879657745361
>>> r = np.empty(N)
>>> timeit.timeit('v.sum(axis=1, out=r)', 'from __main__ import v,r', number=1)
0.09213519096374512

在这里,您清楚地看到numpy.sum在快速运行索引(v是C-连续的)上求和时是最佳的,而在慢速运行轴上求和时则是次优的。有趣的是,对于for循环,相反的模式是正确的:

>>> r = np.zeros(M)
>>> timeit.timeit('for row in v[:]: r += row', 'from __main__ import v,r', number=1)
0.11945700645446777
>>> r = np.zeros(N)
>>> timeit.timeit('for row in v.T[:]: r += row', 'from __main__ import v,r', number=1)
1.2647287845611572

我没有时间检查numpy代码,但我怀疑产生差异的是连续的内存访问或跨步访问。

如这个例子所示,在实现数值算法时,正确的存储器布局具有重要意义。矢量化代码不一定能解决所有问题。