Numba медленный при назначении массиву?

Numba кажется отличным решением для ускорения выполнения числового кода. Однако, когда есть назначения для массива Numba, кажется, медленнее, чем стандартный код Python. Рассмотрим этот пример, сравнивающий четыре альтернативы, с/без Numba, запись в массив/скаляр:

(Расчеты были очень простыми для целей, чтобы сфокусироваться на проблеме, которая является назначением скаляру и присваиванию ячейке массива)

@autojit
def fast_sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

def sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

@autojit
def fast_sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

def sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

Используя IPython% timeit для оценки четырех альтернатив, которые я получил:

In [125]: %timeit fast_sum_arr(arr)
100 loops, best of 3: 10.8 ms per loop

In [126]: %timeit sum_arr(arr)
100 loops, best of 3: 4.11 ms per loop

In [127]: %timeit fast_sum_sclr(arr)
100000 loops, best of 3: 10 us per loop

In [128]: %timeit sum_sclr(arr)
100 loops, best of 3: 2.93 ms per loop

sum_arr, который не был скомпилирован с Numba, более чем в два раза быстрее, чем fast_sum_arr, который был скомпилирован с Numba. С другой стороны, fast_sum_sclr, который был скомпилирован с Numba, более чем на два порядка быстрее, чем sum_sclr, который не был скомпилирован с Numba.

Итак, Numba отлично выполняет задачу ускорения sum_sclr, но на самом деле делает sum_arr медленнее. Единственная разница между sum_sclr и sum_arr заключается в том, что первый присваивает скаляр, а последний присваивает ячейке массива.

Я не знаю, есть ли какое-либо отношение, но я недавно прочитал следующее в блоге http://www.phi-node.com/:

"Оказывается, когда Numba сталкивается с какой-либо конструкцией, она не поддерживает напрямую, она переключается на (очень) медленный путь кода."

Автор блога получил Numba для выполнения намного быстрее, используя инструкцию if вместо Python max().

Любые идеи по этому поводу?

Спасибо,

ФС

Ответ 1

Здесь медленнее функция arr.copy(), а не доступ на запись к массиву. Доказательство:

# -*- coding: utf-8 -*-
from numba import autojit
from Timer import Timer
import numpy as np

@autojit
def fast_sum_arr(arr, z):
    #z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

def sum_arr(arr, z):
    #z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

@autojit
def fast_sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

def sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

if __name__ == '__main__':
    vec1 = np.ones(1000)
    z = vec1.copy()
    with Timer() as t0:
        for i in range(10000):
            pass
    print "time for empty loop ", t0.secs
    print
    with Timer() as t1:
        for i in range(10000):
            sum_arr(vec1, z)
    print "time for sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            fast_sum_arr(vec1, z)
    print "time for fast_sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            sum_sclr(vec1)
    print "time for sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            fast_sum_sclr(vec1)
    print "time for fast_sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6

"""
time for empty loop  0.000312089920044

time for sum_arr       [µs]:    432.02688694
time for fast_sum_arr  [µs]:      7.43598937988
time for sum_arr       [µs]:    284.574580193
time for fast_sum_arr  [µs]:      5.74610233307
"""

Ответ 2

Я не знаю много о numba, но если мы сделаем некоторые основные предположения о том, что он делает под капотом, мы можем сделать вывод, почему версия autojit медленнее и как ускорить ее с незначительными изменениями...

Начнем с sum_arr,

1 def sum_arr(arr):
2     z = arr.copy()
3     M = len(arr)
4     for i in range(M):
5         z[i] += arr[i]
6 
7     return z

Довольно ясно, что здесь происходит, но пусть выбирает строку 5, которую можно переписать как

1 a = arr[i]
2 b = z[i]
3 c = a + b
4 z[i] = c

Python будет продолжать интроперировать это как

1 a = arr.__getitem__(i)
2 b = arr.__getitem__(i) 
3 c = a.__add__(b)
4 z.__setitem__(i, c)

a, b и c - все экземпляры numpy.int64 (или аналогичные)

Я подозреваю, что numba пытается проверить тип даты этих элементов и преобразовать их в некоторые собственные типы данных numba (одно из самых больших замедлений, которое я вижу с помощью numpy-кода, непреднамеренно переключается с типов данных python на numpy-типы данных). Если это действительно так, numba выполняет не менее 3 преобразований, 2 numpy.int64 → native, 1 native → numpy.int64 или, вероятно, хуже с промежуточными (numpy.int64 → python int → native (c INT)). Я подозреваю, что numba добавит дополнительные накладные расходы при проверке типов данных, возможно, вообще не оптимизирует цикл. Посмотрим, что произойдет, если мы удалим изменение типа из цикла...

1 @autojit
2 def fast_sum_arr2(arr):
3     z = arr.tolist()
4     M = len(arr)
5     for i in range(M):
6         z[i] += arr[i]
7 
8     return numpy.array(z)

Тонкое изменение в строке 3, список вместо копии, изменяет тип данных на Python ints, но у нас все еще есть numpy.int64 → native в строке 6. Позвольте переписать это, z [i] + = z [ i]

1 @autojit
2 def fast_sum_arr3(arr):
3     z = arr.tolist()
4     M = len(arr)
5     for i in range(M):
6         z[i] += z[i]
7 
8     return numpy.array(z)

При всех изменениях мы видим довольно существенное ускорение (хотя это не обязательно должно бить чистый python). Конечно, arr + arr, просто глупо быстро.

  1 import numpy
  2 from numba import autojit
  3 
  4 def sum_arr(arr):
  5     z = arr.copy()
  6     M = len(arr)
  7     for i in range(M):
  8         z[i] += arr[i]
  9 
 10     return z
 11 
 12 @autojit
 13 def fast_sum_arr(arr):
 14     z = arr.copy()
 15     M = len(arr)
 16     for i in range(M):
 17         z[i] += arr[i]
 18     
 19     return z
 20 
 21 def sum_arr2(arr):
 22     z = arr.tolist()
 23     M = len(arr)
 24     for i in range(M):
 25         z[i] += arr[i]
 26 
 27     return numpy.array(z)
 28 
 29 @autojit
 30 def fast_sum_arr2(arr):
 31     z = arr.tolist()
 32     M = len(arr)
 33     for i in range(M):
 34         z[i] += arr[i]
 35         
 36     return numpy.array(z)
 37     
 38 def sum_arr3(arr):
 39     z = arr.tolist()
 40     M = len(arr)
 41     for i in range(M):
 42         z[i] += z[i]
 43         
 44     return numpy.array(z)
 45 
 46 @autojit
 47 def fast_sum_arr3(arr):
 48     z = arr.tolist()
 49     M = len(arr)
 50     for i in range(M):
 51         z[i] += z[i]
 52 
 53     return numpy.array(z)
 54 
 55 def sum_arr4(arr):
 56     return arr+arr
 57 
 58 @autojit
 59 def fast_sum_arr4(arr):
 60     return arr+arr
 61 
 62 arr = numpy.arange(1000)

И тайминги

In [1]: %timeit sum_arr(arr)
10000 loops, best of 3: 129 us per loop

In [2]: %timeit sum_arr2(arr)
1000 loops, best of 3: 232 us per loop

In [3]: %timeit sum_arr3(arr)
10000 loops, best of 3: 51.8 us per loop

In [4]: %timeit sum_arr4(arr)
100000 loops, best of 3: 3.68 us per loop

In [5]: %timeit fast_sum_arr(arr)
1000 loops, best of 3: 216 us per loop

In [6]: %timeit fast_sum_arr2(arr)
10000 loops, best of 3: 65.6 us per loop

In [7]: %timeit fast_sum_arr3(arr)
10000 loops, best of 3: 56.5 us per loop

In [8]: %timeit fast_sum_arr4(arr)
100000 loops, best of 3: 2.03 us per loop

Ответ 3

Да, Numba использует ленивую инициализацию, поэтому быстрее второй раз вы его вызываете. При больших массивах, несмотря на ленивую инициализацию, numba все же лучше, чем no-numba.

Попробуйте выполнить следующие раскомментирования разных b

import time
import numpy as np

from numba import jit, autojit


@autojit
def fast_sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

def sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

@autojit
def fast_sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

def sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

b = np.arange(100)
# b = np.arange(1000000)
# b = np.arange(100000000)

print('Vector of len {}\n'.format(len(b)))

print('Sum ARR:\n')

time1 = time.time()
sum_arr(b)
time2 = time.time()
print('No numba:          {}'.format(time2 - time1))

time1 = time.time()
fast_sum_arr(b)
time2 = time.time()
print('Numba first time:  {}'.format(time2 - time1))

time1 = time.time()
fast_sum_arr(b)
time2 = time.time()
print('Numba second time: {}'.format(time2 - time1))

print('\nSum SCLR:\n')

time1 = time.time()
sum_sclr(b)
time2 = time.time()
print('No numba:          {}'.format(time2 - time1))

time1 = time.time()
fast_sum_sclr(b)
time2 = time.time()
print('Numba first time:  {}'.format(time2 - time1))

time1 = time.time()
fast_sum_sclr(b)
time2 = time.time()
print('Numba second time: {}'.format(time2 - time1))

В моей системе с python 3, numba 0.34.0 он получает

"""
Vector of len 100

Sum ARR:

No numba:          7.414817810058594e-05
Numba first time:  0.07130813598632812
Numba second time: 3.814697265625e-06

Sum SCLR:

No numba:          2.6941299438476562e-05
Numba first time:  0.05761408805847168
Numba second time: 1.4066696166992188e-05
"""

"""
Vector of len 1000000

Sum ARR:

No numba:          0.3144559860229492
Numba first time:  0.07181787490844727
Numba second time: 0.0014197826385498047

Sum SCLR:

No numba:          0.15929198265075684
Numba first time:  0.05956888198852539
Numba second time: 0.00037789344787597656
"""

"""
Vector of len 100000000

Sum ARR:

No numba:          30.345629930496216
Numba first time:  0.7232880592346191
Numba second time: 0.586756706237793

Sum SCLR:

No numba:          16.271318912506104
Numba first time:  0.11036324501037598
Numba second time: 0.06010794639587402
"""

Интересно видеть, что разница в вычислительном времени между первым вызовом и вторым уменьшает увеличение размера массива. Я не знаю точно, почему это работает так.