NumPy ufunc 在一个轴上比另一个轴快 2 倍答案

【问题标题】：NumPy ufuncs are 2x faster in one axis over the otherNumPy ufunc 在一个轴上比另一个轴快 2 倍
【发布时间】：2018-07-06 09:54:13
【问题描述】：

我正在做一些计算，并测量了 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 上的速度比cumsum 在轴 0 上的速度几乎 2 倍。为什么会这样？幕后发生了什么？很高兴能清楚地了解其背后的原因。谢谢！

更新：经过一番研究，我意识到如果有人正在构建一个应用程序，他们总是sum只在某个轴上，那么数组应该初始化为适当的顺序：即 C-order 用于轴 = 1 总和或 Fortran 顺序 用于轴 = 0 总和，以节省 CPU 时间。

另外：difference between contiguous and non-contiguous arrays 上的这个出色回答帮助很大！

【问题讨论】：

我的机器上的对比度更大。我可以想象逐行求和对缓存更友好。
@cᴏʟᴅsᴘᴇᴇᴅ 很可能，因为我在集群上尝试这个:)。此外，不仅sum 几乎所有可以在轴上减少的 ufunc 的行为方式都相同

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

【解决方案1】：

数组是row-major。因此，当您对轴 1 求和时，这些数字位于连续的内存阵列中。这允许更好的缓存性能，因此更快的内存访问（参见“Locality of reference”）。我想这就是您在这里看到的效果。

【讨论】：

【解决方案2】：

确实，性能将取决于数组在内存中的顺序：

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)

【讨论】：

【解决方案3】：

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

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

【讨论】：