python - 为什么 numpy 比 for 循环慢

Question

更新:此功能现在位于 sciPy.stats.qmc.discrepancy 中，并已移植到 Cython 并进行了并行处理。

---

我有一个使用一些 for 循环的函数，我想使用 numpy 提高速度。但这似乎并不能解决问题，因为颠簸版本似乎慢了 2 倍。这是代码:

import numpy as np
import itertools
import timeit

def func():
    sample = np.random.random_sample((100, 2))
    
    disc1 = 0
    disc2 = 0
    n_sample = len(sample)
    dim = sample.shape[1]

    for i in range(n_sample):
        prod = 1
        for k in range(dim):
            sub = np.abs(sample[i, k] - 0.5)
            prod *= 1 + 0.5 * sub - 0.5 * sub ** 2
    
        disc1 += prod

    for i, j in itertools.product(range(n_sample), range(n_sample)):
        prod = 1
        for k in range(dim):
            a = 0.5 * np.abs(sample[i, k] - 0.5)
            b = 0.5 * np.abs(sample[j, k] - 0.5)
            c = 0.5 * np.abs(sample[i, k] - sample[j, k])
            prod *= 1 + a + b - c
        disc2 += prod

    c2 = (13 / 12) ** dim - 2 / n_sample * disc1 + 1 / (n_sample ** 2) * disc2


def func_numpy():
    sample = np.random.random_sample((100, 2))

    disc1 = 0
    disc2 = 0
    n_sample = len(sample)
    dim = sample.shape[1]

    disc1 = np.sum(np.prod(1 + 0.5 * np.abs(sample - 0.5) - 0.5 * np.abs(sample - 0.5) ** 2, axis=1))
    
    for i, j in itertools.product(range(n_sample), range(n_sample)):
        disc2 += np.prod(1 + 0.5 * np.abs(sample[i] - 0.5) + 0.5 * np.abs(sample[j] - 0.5) - 0.5 * np.abs(sample[i] - sample[j]))
    
    c2 = (13 / 12) ** dim - 2 / n_sample * disc1 + 1 / (n_sample ** 2) * disc2


print('Normal function time: ' , timeit.repeat('func()', number=20, repeat=5, setup="from __main__ import func"))
print('numpy function time: ', timeit.repeat('func_numpy()', number=20, repeat=5, setup="from __main__ import func_numpy"))

定时输出为:

Normal function time:  [2.831496894999873, 2.832342429959681, 2.8009242500411347, 2.8075121529982425, 2.824807019031141]
numpy function time:  [5.154757721000351, 5.2011515340418555, 5.148996959964279, 5.095560318033677, 5.125199959962629]

我在这里错过了什么？我知道瓶颈是 itertools 部分，因为我有一个 100x100x2 循环而不是之前的 100x2 循环。 您是否看到另一种方法可以做到这一点？

Answer 1

使用 NumPy，人们必须寻求向量化事物，我们当然可以在这里这样做。

仔细观察循环部分，我们在循环启动时沿输入数据 samples 的第一个轴迭代两次:

for i, j in itertools.product(range(n_sample), range(n_sample)):

一旦我们让 broadcasting 可以将这些迭代转换为向量化操作处理那些。

现在，要获得完全矢量化的解决方案，我们需要更多的内存空间，特别是 (N,N,M)，其中 (N,M)是输入数据的形状。

这里另一个值得注意的方面是，在每次迭代中，我们没有做很多工作，因为我们对每一行执行操作，并且每一行仅包含给定样本的 2 元素。因此，出现的想法是我们可以沿着 M 运行一个循环，这样在每次迭代中，我们都会计算 prod 并累加。因此，对于给定的样本，它只有两个循环迭代。

跳出循环，我们将得到累积的prod，它只需要对disc2求和作为最终输出。

下面是实现上述想法的实现 -

prod_arr = 1
for i in range(sample.shape[1]):
    si = sample[:,i]
    prod_arr *= 1 + 0.5 * np.abs(si[:,None] - 0.5) + 0.5 * np.abs(si - 0.5) - \
                                    0.5 * np.abs(si[:,None] - si)
disc2 = prod_arr.sum()

运行时测试

下面列出了原始方法中循环部分的精简版本和作为方法的修改版本:

def org_app(sample):
    disc2 = 0
    n_sample = len(sample)
    for i, j in itertools.product(range(n_sample), range(n_sample)):
        disc2 += np.prod(1 + 0.5 * np.abs(sample[i] - 0.5) + 0.5 * \
            np.abs(sample[j] - 0.5) - 0.5 * np.abs(sample[i] - sample[j]))
    return disc2


def mod_app(sample):
    prod_arr = 1
    for i in range(sample.shape[1]):
        si = sample[:,i]
        prod_arr *= 1 + 0.5 * np.abs(si[:,None] - 0.5) + 0.5 * np.abs(si - 0.5) - \
                                        0.5 * np.abs(si[:,None] - si)
    disc2 = prod_arr.sum()
    return disc2

时间和验证-

In [10]: sample = np.random.random_sample((100, 2))

In [11]: org_app(sample)
Out[11]: 11934.878683659041

In [12]: mod_app(sample)
Out[12]: 11934.878683659068

In [14]: %timeit org_app(sample)
10 loops, best of 3: 84.4 ms per loop

In [15]: %timeit mod_app(sample)
10000 loops, best of 3: 94.6 µs per loop

关于 900x 加速!好吧，这应该足以激励我们，希望尽可能地对事物进行矢量化。