是什么让 numpy.sum 和 numpy.cumsum 之间的表现不同？

Question

如果我的理解是正确的，np.sum 和np.cumsum 都需要O(n) 时间。但是，当我按顺序（在不同的轴上）对同一个矩阵执行这两种操作时，似乎np.sum 和np.cumsum 的顺序使整体性能大不相同，尽管结果与预期相同。

如果我先对每一行执行np.cumsum沿列方向（axis=1），然后对所有行执行np.sum（axis=0），则需要更长的时间。

如果我先沿行执行np.sum (axis=0)，然后对一维数组执行np.cumsum，所需时间会更短。

我的假设是np.cumsum 将花费更多时间进行数据分配/操作，因为它会产生比np.sum 更多的数据，因此如果有更多np.cumsum 的操作将花费更长的时间。

这是我的测试代码和结果

import numpy as np
import time

b = np.zeros((1000, 1000))
for i in range(1000):
    b[i] = np.array(range(1000))

time_start = time.time()
for i in range(1000):
    c = np.cumsum(b, axis=1)
    d = np.sum(c, axis=0)
time_end = time.time()
print(f"np.sum(np.cumsum(...)) time: {time_end - time_start}")


time_start = time.time()
for i in range(1000):
    c = np.cumsum(np.sum(b, axis=0))
time_end = time.time()
print(f"np.cumsum(np.sum(...)) time: {time_end - time_start}")

结果是

np.sum(np.cumsum(...)) time: 3.6612446308135986
np.cumsum(np.sum(...)) time: 0.38796162605285645

Answer 1

如果我的理解是正确的，np.sum 和 np.cumsum 都需要 O(n) 时间。

正确（假设输入相同）！

看来 np.sum 和 np.cumsum 的顺序让整体表现大不一样

np.cumsum 在相同的输入 上比np.sum 更昂贵，因为np.cumsum 需要分配和写入一个输出数组（可能非常大） .此外，假设浮点运算是关联的，np.sum 的实现很容易优化，而np.cumsum 的实现很难优化。

我的假设是 np.cumsum 在数据分配/操作上会花费更多时间，因为它会产生比 np.sum 更多的数据，因此如果 np.cumsum 的操作更多，则需要更长的时间。

问题是第一个实现比第二个实现要多得多。这主要是它变慢的原因，特别是因为内存读/写。事实上，np.cumsum 在第一个实现中产生了一个大的二维数组，必须由np.sum 计算。 写入和读取这个临时数组很昂贵。在示例中，第一个实现需要将 14.9 GiB 从/移入内存层次结构（可能在 RAM 中），仅用于临时数组。请注意，构建一个临时数组也会降低计算缓存友好，因为它需要更多空间，因此 b 和 c 可能不再适合缓存（b 和 @ 987654331@ 在我的机器上每个占用 8 MiB，而我的 CPU 有一个 9 MiB 的 L3 缓存，这意味着不是两者都可以放入缓存中），而不是第二种实现。 RAM 的吞吐量通常远小于 CPU 缓存之一。

请注意，第二个实现也会产生临时数组。这个版本应该分配与第一个一样多的临时数组。但是，第二个实现会生成1000 倍的临时数组！因此，在此版本中，对np.cumsum 的调用要快得多。在我的机器上，b 主要存储在快速 L3 缓存中，而临时数组则存储在非常快速的 L1 缓存中。

Answer 2

查看应用于小数组的替代方案。

In [45]: b = np.arange(24).reshape(4,6)
In [46]: b
Out[46]: 
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23]])

sum(cumsum(b)):

In [47]: np.cumsum(b, axis=1)         # same shape as b
Out[47]: 
array([[  0,   1,   3,   6,  10,  15],
       [  6,  13,  21,  30,  40,  51],
       [ 12,  25,  39,  54,  70,  87],
       [ 18,  37,  57,  78, 100, 123]])
In [48]: np.sum(_, axis=0)
Out[48]: array([ 36,  76, 120, 168, 220, 276])

先求和：

In [49]: np.sum(b, axis=0)            # much reduced array
Out[49]: array([36, 40, 44, 48, 52, 56])
In [50]: np.cumsum(_)
Out[50]: array([ 36,  76, 120, 168, 220, 276])

虽然第一步对相同数量的元素（整个 b）进行操作，但当首先完成 sum 时，第二步将使用大大减少的数组。

实际上对于这个小样本，时间差异很小：

In [57]: timeit np.sum(np.cumsum(b,axis=1), axis=0)
17.8 µs ± 29.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [58]: timeit np.cumsum(np.sum(b, axis=0))
16.1 µs ± 30.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

但是对于更大的阵列，由于减少而节省的时间变得非常重要：

In [59]: b = np.random.randint(0,1000,(1000,1000))
In [60]: timeit np.sum(np.cumsum(b,axis=1), axis=0)
5.11 ms ± 9.99 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [61]: timeit np.cumsum(np.sum(b, axis=0))
1.64 ms ± 649 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

单独看第一步：

In [62]: timeit np.cumsum(b,axis=1)
3.73 ms ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [63]: timeit np.sum(b, axis=0)
1.62 ms ± 1.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

cumsum，因为它返回完整的b.shape 数组比较慢。

O(n) 告诉我们有关操作如何扩展的一些信息，例如在小型阵列与大型阵列上进行操作时。但它不能用于将一项操作与另一项操作进行比较。除了numpy，核心计算在编译后的代码中可以快速完成，但内存管理和其他设置问题可能会占用核心时间。

编辑

np.sum本质上是np.add.reduce：

In [79]: timeit np.add.reduce(b,axis=0)
1.59 ms ± 4.73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [80]: timeit np.sum(b,axis=0)
1.64 ms ± 5.48 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

而np.cumsum 是np.add.accumulate：

In [81]: timeit np.cumsum(b,axis=0)
10.2 ms ± 357 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [82]: timeit np.add.accumulate(b,axis=0)
10 ms ± 14.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)