2D 数组上的 numpy.where 比 1D 数组上 numpy.where 的列表理解慢答案

【问题标题】：numpy.where on 2D array is slower than list comprehension of numpy.where on 1D array2D 数组上的 numpy.where 比 1D 数组上 numpy.where 的列表理解慢
【发布时间】：2022-01-12 06:04:33
【问题描述】：

为什么在这种情况下 Numpy 比列表推导慢？
矢量化此网格结构的最佳方法是什么？

In [1]: import numpy as np

In [2]: mesh = np.linspace(-1, 1, 3000)

In [3]: rowwise, colwise = np.meshgrid(mesh, mesh)

In [4]: f = lambda x, y: np.where(x > y, x**2, x**3)

# Using 2D arrays:
In [5]: %timeit f(colwise, rowwise)
285 ms ± 2.25 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Using 1D array and list-comprehension:
In [6]: %timeit np.array([f(x, mesh) for x in mesh])
58 ms ± 2.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Equivalent result
In [7]: np.allclose(f(colwise, rowwise), np.array([f(x, mesh) for x in mesh]))
True

【问题讨论】：

您在这两种情况下都使用where。但是对于 [5]，所有三个参数都是 (3000,3000) 形的。在 [6] 中，虽然 f 被调用了 3000 次，但 x 参数是一个标量。使用f(mesh,x) 会减慢它的速度。另一方面，在meshgrid 中使用sparse=True 会加快[5]，因为x arg 仅为(3000,)。
where 是一个函数。它的参数在传递给它之前会被完整评估。 where` 然后返回一个与三个数组的 broadcasted 并集兼容的新数组。不要将其视为迭代器。

标签： python arrays numpy performance

【解决方案1】：

为什么在这种情况下 Numpy 比列表推导慢？

你基本上遇到了两个问题。

第一个是缓存利用率，第二个版本使用空间的较小子集 ~~(3000,1)~~ (1,3000) 进行计算，这可以很好地适合您的缓存，因此x>y, x**2 , x**3 可以全部适合您的缓存，这有点加快速度，第一个版本是为 3000x3000 数组（900 万个条目）计算这 3 个中的每一个，该数组永远不会位于您的缓存中（通常约为 2-5 MB），然后调用 np.where 必须从您的数据中获取部分数据RAM（而不是缓存）以进行内存复制，然后将其逐个返回到您的 RAM，这非常昂贵。

np.where 的 numpy 实现也有点不知道对齐，并且按列访问数组，而不是按行访问，因此它本质上是从 RAM 中获取每个条目，根本不使用缓存。

你的列表理解实际上解决了这个问题，因为它只在给定时间处理一小部分数据，因此所有数据都可以放在你的缓存中，但它仍然使用 np.where，它只是强制它使用逐行访问，从而利用您的缓存。

第二个问题是x**2和x**3的计算，这是一个浮点取幂，很贵，考虑换成x*x和x*x*x

矢量化这个网格结构的最佳方法是什么？

显然你是用第二种方法写的。

利用缓存进行更快但不必要的优化是用 C 语言编写自己的代码并在 python 中调用它，这样你就不必评估x*x or x*x*x，除非你需要，也不必存储 x>y,x*x,x*x*x 但加速不值得麻烦。

【讨论】：

缓存是否解释了使用 (3000,1) 和 (1,3000) broadcastable 数组的更快速度？
这取决于上下文和使用，有时 numpy 会在应用手头的操作之前将它们都广播为 (3000,3000)，因此不会有加速，但一般来说你应该期望只要您没有被 python 解释器拖慢，就可以使用较小的数组来完成较大数组的工作。
在我在 (3000,3000) 数组上测试 x**3 时，这是主要问题，比 O(n) 缩放稍差。

【解决方案2】：

In [1]: In [2]: mesh = np.linspace(-1, 1, 3000)
   ...: In [3]: rowwise, colwise = np.meshgrid(mesh, mesh)
   ...: In [4]: f = lambda x, y: np.where(x > y, x**2, x**3)

另外让我们制作稀疏网格：

In [2]: r1,c1 = np.meshgrid(mesh,mesh,sparse=True)
In [3]: rowwise.shape
Out[3]: (3000, 3000)
In [4]: r1.shape
Out[4]: (1, 3000)

使用稀疏网格，时间甚至比您的迭代更好：

In [5]: timeit f(colwise, rowwise)
645 ms ± 57.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [6]: timeit f(c1,r1)
108 ms ± 3.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [7]: timeit np.array([f(x, mesh) for x in mesh])
166 ms ± 13.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

另一个答案强调缓存。其他帖子表明，适度的迭代可以比处理非常大的数组更快，例如在使用matmul 时。我不知道是缓存还是其他一些内存管理复杂性会减慢速度。

但在3000*3000*8 bytes 我不确定这是这里的问题。相反，我认为现在是 x**2 和 x**3 表达式需要的时候了。

where 的参数在传入之前进行评估。

条件表达式需要一些时间：

In [8]: timeit colwise>rowwise
24.2 ms ± 71.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

但是 (3000,3000) 数组的幂表达式占用了总时间的大部分：

In [9]: timeit rowwise**3
467 ms ± 8.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

将其与稀疏等效项所需的时间进行对比：

In [10]: timeit r1**3
142 µs ± 150 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

这次快了 3288 倍；这比 O(n) 缩放要差一些。

重复乘法更好：

In [11]: timeit rowwise*rowwise*rowwise
116 ms ± 12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

在[f(x, mesh) for x in mesh] 中，x**3 在标量上运行，因此速度很快，即使重复了 3000 次。

其实如果我们把功率计算从时序中去掉，整个数组where还是比较快的：

In [15]: %%timeit x2,x3 = rowwise**2, rowwise**3
    ...: np.where(rowwise>colwise, x2,x3)
89.8 ms ± 3.99 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

【讨论】：