比较 Numpy 和 Matlab 数组求和速度

Question

我最近使用 Numpy 将 MATLAB 脚本转换为 Python，发现它的运行速度明显变慢。我期望类似的性能，所以我想知道我是否做错了什么。

作为简化示例，我手动对几何级数求和：

MATLAB 版本：

function s = array_sum(a, array_size, iterations)
    s = zeros(array_size);
    for m = 1:iterations
        s = a + 0.5*s;
    end
end

% benchmark code
array_size = 500
iterations = 500
a = randn(array_size)
f = @() array_sum(a, array_size, iterations);
fprintf('run time: %.2f ms\n', timeit(f)*1e3);

Python/Numpy 版本：

import numpy as np
import timeit

def array_sum(a, array_size, iterations):
    s = np.zeros((array_size, array_size))
    for m in range(iterations):
        s = a + 0.5*s
    return s

array_size = 500
iterations = 500
a = np.random.randn(array_size, array_size)
timeit_iterations = 10
t1 = timeit.timeit(lambda: array_sum(a, array_size, iterations),
                   number=timeit_iterations)
print("run time: {:.2f} ms".format(1e3*t1/timeit_iterations))

在我的机器上，MATLAB 在 58 毫秒内完成。 Python 版本的运行时间为 292 毫秒，或者慢了 5 倍。

我还尝试通过添加 Numba JIT 装饰器 @jit('f8[:,:](i8, i8)', nopython=True) 来加速 Python 代码，但时间仅下降到 236 毫秒（慢 4 倍）。

这比我预期的要慢。我是否不正确地使用 timeit？我的 Python 代码有问题吗？

编辑：编辑以便在基准函数之外创建随机矩阵。

编辑 2：我使用 Torch 而不是 Numpy 运行基准测试（计算总和为 s = torch.add(s, 0.5, a)），它在我的计算机上运行仅 52 毫秒！

Answer 1

根据我的经验，在使用 numba 的 jit 函数时，将数组操作扩展为循环通常更快。所以我尝试将你的python函数重写为：

@jit(nopython=True, cache=True)
def array_sum_numba(a, array_size, iterations):
    s = np.zeros((array_size, array_size))
    for m in range(iterations):
        for i in range(array_size):
            for j in range(array_size):
                s[i,j] = a[i,j] + 0.5 * s[i,j]
    return s

出于好奇，我还测试了@percusse 的版本，对参数稍作修改：

def array_sum2(r, array_size, iterations):
    s = np.zeros((array_size, array_size))
    for m in range(iterations):
        s /= 2
        s += r
    return s

我机器上的测试结果是：

• 原版运行时间：143.83 ms
• numba jited loop 版本运行时间：26.99 毫秒
• @percusse 的版本运行时间：61.38 毫秒

这个结果在我的预期之内。值得一提的是，我已将 timeit 迭代次数增加到 50，这大大减少了 numba 版本的时间。

总结：如果使用 numba 的 jit 并在循环中编写函数，Python 代码仍然可以显着加速。我的机器上没有要测试的 Matlab，但我的猜测是使用 numba，python 版本更快。

Answer 2

由于您正在更新适用于就地操作的相同变量，因此您可以将函数更新为

def array_sum2(array_size, iterations):
    s = np.zeros((array_size, array_size))
    r = np.random.randn(array_size, array_size)
    for m in range(iterations):
        s /= 2
        s += r
    return s

与array_sum相比，这在我的机器上带来了以下速度优势

run time: 157.32 ms
run time2: 672.43 ms

Answer 3

时间包括randn 调用以及总和：

In [68]: timeit array_sum(array_size, 0)
16.6 ms ± 436 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [69]: timeit array_sum(array_size, 1)
18.9 ms ± 293 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [70]: timeit array_sum(array_size, 20)
55.5 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [71]: (55-16)/20
Out[71]: 1.95

所以设置需要 16 毫秒，每次迭代需要 2 毫秒。 500 次迭代的相同模式。

MATLAB 进行一些 JIT 编译。我不知道这里是不是这样。我没有 MATLAB 来测试。八度音阶（无timeit）

>> t = time(); array_sum(500,0); (time()-t)*1000
ans =  13.704
>> t = time(); array_sum(500,1); (time()-t)*1000
ans =  16.219
>> t = time(); array_sum(500,20); (time()-t)*1000
ans =  82.346
>> t = time(); array_sum(500,500); (time()-t)*1000
ans =  1610.6

Octave 的random 更快，但每次迭代的总和更慢。