NumPy的ufunc在某个轴上比另一个轴快两倍。

Question

我正在进行一些计算，并测量了不同轴上的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)

cumsum沿着轴1的速度几乎比沿着轴0的cumsum快2倍。为什么会这样，背后发生了什么？有一个清晰的理解原因将是很好的。谢谢！

---

更新：经过一番研究，我意识到如果有人正在构建一个应用程序，在该应用程序中他们总是对特定轴进行求和，则应以适当的顺序初始化数组：即对于轴=1的求和使用C-order，对于轴=0的求和使用Fortran-order，以节省CPU时间。

此外：这篇关于连续和非连续数组之间差异的优秀答案帮了很大忙！

- kmario23

---

在我的电脑上，对比度更大。我能想象按行求和更加缓存友好。 - cs95

@cᴏʟᴅsᴘᴇᴇᴅ 可能确实是这样，因为我是在集群上尝试的 :). 而且，不仅是 sum，几乎所有可以在轴上进行缩减的 ufunc 都表现出相同的行为。 - kmario23

@kmario23 请不要包含问题的解决方案（请另发一条回答）。 - undefined

Answer 1

您有一个正方形的数组，长下面这样：

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字节（使用NumPy的默认dtype为8字节浮点数，可以装载8个值）。

你可以在构建数组时使用order参数来控制NumPy使用哪种布局。

了解更多信息，请参见：https://en.wikipedia.org/wiki/Row-_and_column-major_order

- John Zwinck

Answer 2

数组是行主序的。因此，在对轴1求和时，数字在连续的内存数组中被找到。这可以提高缓存性能，从而实现更快的内存访问（参见“引用局部性”）。我认为这就是你在这里看到的效果。

- rain city

Answer 3

事实上，性能将取决于内存中数组的顺序:

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

- AGN Gazer