Numba：如何以编程方式打开 on/off 即时编译（设置 NUMBA_DISABLE_JIT 环境变量）？

Question

我编写了一个脚本，旨在测试我添加到多个函数的 @numba.jit 装饰器的功效。我想先在没有即时编译的情况下运行带注释的函数，然后使用 @numba.jit 装饰器再次运行，以便比较两个结果。

我已经尝试通过 os.environ 修改 NUMBA_DISABLE_JIT 环境变量的值来做到这一点，但我不确定这是否会产生预期的效果。例如：

# 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是否真的如我假设的disabling/enabling所有装饰有@numba.jit的函数的JIT编译效果？如果没有那么有没有更好的方法给这只猫剥皮？

Answer 1

我认为该标志只对 expensive_function 的第一次调用有影响，所以没有按照您的意愿进行。

使用 numba，您始终可以使用 .py_func 访问原始 python 函数，因此这可能是一种更简单的方法

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)

Answer 2

从 numba 导入配置并将其设置为 True 以禁用 @jit 和 False 以启用@jit。

from numba import jit, config

config.DISABLE_JIT = True

Numba：如何以编程方式打开 on/off 即时编译（设置 NUMBA_DISABLE_JIT 环境变量）？

Numba: How to turn on/off just-in-time compilation programmatically (setting NUMBA_DISABLE_JIT environment variable)?

python

jit

numba