NumPy ufuncs在一个轴上比另一个轴快2倍

时间:2018-01-27 01:53:18

标签: python performance numpy numpy-ufunc numpy-ndarray

我正在做一些计算,并在不同的轴上测量 ufuncs 的性能,如np.cumsum,以使代码更高效。

In [51]: arr = np.arange(int(1E6)).reshape(int(1E3), -1)

In [52]: %timeit arr.cumsum(axis=1)
2.27 ms ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [53]: %timeit arr.cumsum(axis=0)
4.16 ms ± 10.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
轴1上的

cumsum在轴0上比cumsum >快2倍为什么会这样,在幕后发生了什么?能够清楚地了解其背后的原因是很好的。谢谢!

更新:经过一番研究后,我意识到,如果某人正在构建一个只在某个轴上总和的应用程序,那么该数组应该在适当的顺序:即 C-order for axis = 1 sums或 Fortran-order for axis = 0 sums,以节省CPU时间。

另外:difference between contiguous and non-contiguous arrays的这个优秀答案帮助了很多!

3 个答案:

答案 0 :(得分:8)

你有一个方阵。它看起来像这样:

1 2 3
4 5 6
7 8 9

但是计算机内存是线性寻址的,所以计算机看起来像这样:

1 2 3 4 5 6 7 8 9

或者,如果您考虑一下,它可能看起来像这样:

1 4 7 2 5 8 3 6 9

如果您要对[1 2 3][4 5 6](一行)求和,第一个布局会更快。如果您要对[1 4 7][2 5 8]求和,则第二种布局会更快。

这是因为从内存加载数据发生了一个"缓存行"一次,通常是64个字节(8个值,NumPy'默认dtype为8个字节的浮点数)。

您可以使用order参数控制NumPy在构造数组时使用的布局。

有关详情,请参阅:https://en.wikipedia.org/wiki/Row-_and_column-major_order

答案 1 :(得分:7)

数组是row-major。因此,当你在轴1上求和时, 数字在连续的内存数组中找到。这样可以提高缓存性能,从而加快内存访问速度(参见“Locality of reference”)。我认为这就是你在这里看到的效果。

答案 2 :(得分:1)

实际上,性能取决于内存中数组的顺序:

In [36]: arr = np.arange(int(1E6)).reshape(int(1E3), -1)

In [37]: arrf = np.asfortranarray(arr) # change order

In [38]: %timeit arr.cumsum(axis=1)
1.99 ms ± 32.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [39]: %timeit arr.cumsum(axis=0)
14.6 ms ± 229 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [41]: %timeit arrf.cumsum(axis=0)
1.96 ms ± 19.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [42]: %timeit arrf.cumsum(axis=1)
14.6 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

有关详细信息,请参阅https://docs.scipy.org/doc/numpy-1.13.0/reference/internals.html#multidimensional-array-indexing-order-issues