python - 比较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 毫秒
• numba jitted 循环版本运行时间:26.99 毫秒
• @percusse 的版本运行时间:61.38 毫秒

这个结果在我的意料之中。值得一提的是，我已将 timeit 迭代次数增加到 50，这导致 numba 版本的时间显着减少。

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