Question

我有一个项目，我计算许多大型数组的累积总和。在我的服务器上2秒这一步是一个很大的瓶颈。有什么办法可以加快速度吗？注意，这些阵列代表温度测量。所以它们是浮点值，既可以是负数也可以是正数。虽然我有更多的内核可用，但我已经在其他地方使用并行处理，所以在这种情况下不会加快速度。

import numpy as np
import time

forcing = np.random.rand(380*1400*620).reshape((380,1400,620))

start = time.time()
forcing.cumsum(axis=0)
np_time = time.time() - start
print(np_time)
2.085033416748047

Answer 1

作为mentioned by daniel451，numpy未并行化cumsum操作，因此您可以明确地将其并行化以获得至少一点性能。

例如，使用multiprocessing.dummy（multiprocessing API的线程支持副本），您可以执行以下操作：

import numpy as np
from multiprocessing.dummy import Pool
from itertools import repeat

forcing = np.random.rand(380*1400*620).reshape((380,1400,620))

# Make an output array of matching size, that can be populated piecemeal
# in each thread
forceres = np.zeros_like(forcing)

# Compute cumsum in parallel over second dimension
with Pool() as pool:
    # Use module function with np.rollaxis to avoid need to define
    # worker to do slicing
    pool.starmap(np.cumsum, zip(np.rollaxis(forcing, 1), repeat(0), repeat(None), np.rollaxis(forceres, 1)))

我在八核计算机上使用ipython3的{{1}} / %time魔法对此进行了测试，发现它与原始代码的运行时间相比减少了近70％，从5.49秒开始到1.73秒;你的机器显然更快，所以如果你的机器发生同样的加速，我希望它需要大约0.66秒。

我的比较是：

%%time

VS

>>> %%time
... forcesres = np.zeros_like(forcing)
... with Pool() as pool:
...     pool.starmap(np.cumsum, zip(np.rollaxis(forcing, 1), repeat(0), repeat(None), np.rollaxis(forceres, 1)))
CPU times: user 10 s, sys: 213 ms, total: 10.2 s
Wall time: 1.73 s

任何加速numpy.cumsum（）的方法

1 个答案: