Сравнение скорости суммирования массива Numpy и Matlab

Недавно я преобразовал MATLAB script в Python с помощью Numpy и обнаружил, что он работает значительно медленнее. Я ожидал подобную производительность, поэтому мне интересно, что я делаю что-то неправильно.

Как урезанный пример, я вручную суммирую геометрический ряд:

Версия 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 мс, или на 5X медленнее.

Я также попытался ускорить код Python, добавив декоратор Numba JIT @jit('f8[:,:](i8, i8)', nopython=True), но время только упало до 236 мс (на 4 раза медленнее).

Это медленнее, чем я ожидал. Я использую timeit неправильно? Что-то не так с моим кодом Python?

EDIT: отредактировано так, что случайная матрица создается вне контрольной функции.

РЕДАКТИРОВАТЬ 2: Я использовал контрольную точку вместо Torch (вычисляя сумму как s = torch.add(s, 0.5, a)), и она работает всего на 52 мс на моем компьютере!

Ответ 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 loop version время запуска: 26.99 мс
Время работы версии @percusse: 61,38 мс

Этот результат в моих ожиданиях. Стоит отметить, что я увеличил время итераций до 50, что приводит к некоторому значительному сокращению времени для версии numba.

Вкратце: код Python может быть значительно ускорен, если вы используете numba jit и записываете функцию в циклы. У меня нет Matlab на моей машине, чтобы проверить, но я думаю, что с numba версия python выполняется быстрее.

Ответ 2

Поскольку вы обновляете одну и ту же переменную, подходящую для операций inplace, вы можете обновить свою функцию как

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

Ответ 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. В 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 работает быстрее, но каждая итерационная сумма медленнее.