Numba: Как программно включить / отключить своевременную компиляцию (установив переменную среды NUMBA_DISABLE_JIT)? - PullRequest
Новая
       15

Numba: Как программно включить / отключить своевременную компиляцию (установив переменную среды NUMBA_DISABLE_JIT)?

0 голосов
/ 27 декабря 2018

Я написал скрипт, предназначенный для проверки эффективности декораторов @numba.jit, которые я добавил к нескольким функциям.Я хочу сначала выполнить аннотированные функции без компиляции «точно в срок», а затем выполнить это снова с задействованными декораторами @numba.jit, чтобы затем сравнить два результата.

Я пытался выполнитьизменив значение переменной окружения NUMBA_DISABLE_JIT с помощью os.environ, но я пока не уверен, что это дает желаемый эффект.Например: 

# run first without and then with numba's just-in-time compilation
for flag in [1, 0]:

    # enable/disable numba's just-in-time compilation
    os.environ["NUMBA_DISABLE_JIT"] = str(flag)

    # time an arbitrary number of iterations of the JIT decorated function
    start = time.time()
    for _ in range(1000):
        expensive_function()
    end = time.time()

    # display the elapsed time
    if flag == 0:
        preposition = "with"
    else:
        preposition = "without"
    print("Elapsed time " + preposition + " numba: {t}".format(t=(end - start)))

Имеет ли приведенная выше настройка переменной среды DISABLE_NUMBA_JIT фактически отключение / включение JIT-компиляции всех функций, отмеченных @numba.jit, как я полагаю?Если нет, то есть ли лучший способ снять кожу с этой кошки?

1 Ответ

0 голосов
/ 27 декабря 2018

Я думаю, что флаг влияет только на первый вызов expensive_function, поэтому не делает то, что вам нравится.

С помощью numba вы всегда можете получить доступ к исходной функции python с помощью .py_func, так что это может быть более простой способ сделать это 

import numba

@numba.njit
def expensive_function(arr):
    ans = 0.0
    for a in arr:
        ans += a
    return ans

arr = np.random.randn(1_000_000)

In [21]: %timeit expensive_function(arr)
# 1.51 ms ± 24.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit expensive_function.py_func(arr)
# 134 ms ± 11 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
...