numpy.sum 性能取决于轴

Question

对 numpy 数组中的维度求和时，第一个轴和最后一个轴之间是否存在性能差异？

具体来说，考虑下面的代码，sum1和sum2哪个会执行得更快？

import numpy as np
a = np.ones((1000,200))
b = np.ones((200,1000))
sum1 = np.sum(a, axis=0)
sum2 = np.sum(b, axis=-1)

我相信这个问题实际上归结为 numpy 如何在内部存储维度，并且可以将其覆盖以使用逐行或逐列格式。但是，当使用默认设置时，其中哪一个会更快？还有，N维数组呢？

Answer 1

很容易检查是否存在性能差异（IPython，我增加了一点数字以获得更明显的差异）：

import numpy as np

a = np.ones((10000, 2000))
b = np.ones((2000, 10000))
%timeit np.sum(a, axis=0)
# 27.6 ms ± 541 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit  np.sum(b, axis=-1)
# 34.6 ms ± 876 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

现在，当您遇到np.sum 的实际性能问题时，您可能已经内存不足，但是是的，存在差异。默认情况下，NumPy 数组存储在row-major order 中，因此首先进入第一行，然后进入第二行，依此类推。因此，外部维度的求和（或运算）更快，因为缓存将是方式更加有效。简单地说，在第一种情况下，当你得到数组的第一个元素时，一堆连续的数据会随之进入缓存，所以当你想对下一个元素求和时，它们已经在那里了。另一方面，在第二种情况下，要求和的元素彼此相距很远（实际上是 2000 个元素的距离），因此缓存在列方面不会有太大帮助。这并不是说缓存根本没有帮助，因为您正在对所有列求和，所以缓存的数据仍然会在一定程度上被重用，但效率不高。这是一个相当粗略的近似，通常有几个缓存级别，一些在内核之间共享，一些不共享，理解一个或另一个代码对其的确切影响是一个复杂的话题，但总体思路是成立的。