NumPy:“vectorize”的替代方案,让我可以访问数组

NumPy: Alternative to `vectorize` that lets me access the array

我有这个代码:

output_array = np.vectorize(f, otypes='d')(input_array)

我想用这段代码替换它,它应该提供相同的输出:

output_array = np.ndarray(input_array.shape, dtype='d')
for i, item in enumerate(input_array):
    output_array[i] = f(item)

我想要第二个版本的原因是我可以在计算时在单独的线程中开始迭代 output_array。 (是的,我知道 GIL,那部分已经处理好了。)

不幸的是,for 循环非常慢,即使我没有在单独的线程上处理数据也是如此。我在 CPython 和 PyPy3 上对它进行了基准测试,这是我的目标平台。在 CPython 上它比 vectorize 慢 3 倍,而在 PyPy3 上它比 vectorize!

慢 67 倍

尽管 Numpy 文档说“提供 vectorize 函数主要是为了方便,而不是为了性能。实现本质上是一个 for 循环。”

知道为什么我的实施速度很慢吗?如何快速实施并允许我在完成之前使用 output_array

numpy 中的矢量化 通过 for 循环提高数值计算的性能,因为它减少了一些 Python 解释器动态确定一些开销的开销运行时数据的特征(如其类型和位置)。

一个 ndarray 对象可以做很多优化计算速度的强大事情,比如确保它包含的数据在内存中是同质且连续的。

对于 n 维数据,默认情况下,numpy 将值存储在 C ordering (row major) 中,以便一行的连续元素彼此相邻存储。

它也是一个数据块的跨步视图,这意味着当你初始化它的时候,你也在说里面的每个项目到底占用了多少字节的内存-- 换句话说,你需要迈出多大的一步才能迈出下一步。

当您在 ndarray 对象上使用矢量化函数时,它的行为更像是 C 代码而不是 Python。也就是说,它直接对内存中的值进行操作并修改它们,并利用所有存储和类型优化。

我怀疑当您不向量化函数 f 时,您会增加很多 Python 解释器的开销。成功矢量化后,该函数的大部分将在 C 中执行,并避免 Python 造成的大部分减速。我想知道 enumerate 是否未能像 numpy 的 nditer 对象那样利用底层数据结构,但我对 nditer、numpy ufuncs 和显式循环进行了基准测试 enumerate,而 enumerate 实际上是最快的迭代器,所以我猜您的自定义函数可能是时间的罪魁祸首。这是有道理的,尤其是考虑到 PyPy 的减速

示例基准:

a = np.ndarray()

>>> %%time
>>> for x in np.nditer(a, flags=['external_loop']):
>>> ....    x*x
CPU times: user 201 ms, sys: 219 ms, total: 420 ms
Wall time: 420 ms

>>> %%time
>>> np.square(a)
CPU times: user 201 ms, sys: 180 ms, total: 381 ms
Wall time: 380 ms

>>> %%time
>>> for i, x in enumerate(a):
>>> ....    x*x

CPU times: user 78.5 ms, sys: 1.79 ms, total: 80.3 ms
Wall time: 79.4 ms

Sebastian Berg 给了我一个解决方案。当迭代输入数组中的项目时,使用 item.item() 而不仅仅是 item。这将 numpy.float64 对象变为正常的 Python 浮点数,使一切变得更快并解决了我的特定问题:)