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 jitted loop version运行时间：26.99 ms
@ percusse的版本运行时间：61.38 ms

这个结果符合我的期望。值得一提的是，我将时间迭代次数增加到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

因此设置为16ms，每次迭代为2ms。具有500次迭代的相同模式。

MATLAB做了一些JIT编译。我不知道这是不是这样的。我没有MATLAB来测试。在Octave（没有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速度更快，但每次迭代总和更慢。

比较Numpy和Matlab数组求和速度

3 个答案: