已解决:在 Numba 中使用全局变量的替代方法

Resolved: Alternative to using Global Variable in Numba

我目前正在使用 scipy 最小化来解决我的优化问题,但计算时间很长。我遇到了 numba,它可以用来减少计算时间。但是当我尝试在我的 objective 函数上使用它时,它会抛出以下错误。

TypingError: Failed in nopython mode pipeline (step: ensure IR is legal prior to lowering) The use of a reflected list(int64)<iv=None> type, assigned to variable 'wInt' in globals, is not supported as globals are considered compile-time constants and there is no known way to compile a reflected list(int64)<iv=None> type as a constant.

这是我目前在 objective 函数中使用的示例代码。

#x is a list returned by a function and is run only once at the 
# -beginning of the code execution.
x = someFunc()

@jit(float64(int64), nopython=True, parallel=True)
def fast_rosenbrock(N):
    out = 0.0
    for i in range(N-1):
        out += 100.0 * (x[i+1] - x[i]**2)**2 / (1 - x[i])**2
    return out

objective函数使用了一个全局变量,该变量是通过调用函数获得的。我担心如果我将其设为本地,则会重复计算相应的值,我想避免这种情况,因为函数非常大并且只需要 运行 一次。我该如何解决?

编辑 1:

尝试将 x 作为参数传递。它可以在没有 numba 的情况下工作,但是当我将它变成一个 jitted 函数时,它会抛出一个错误。

没有 numba,我得到了想要的结果:

def fast_rosenbrock(x, N):
    out = 0.0
    for i in range(N-1):
        out += 100.0 * (x[i+1] - x[i]**2)**2 / (1 - x[i])**2
    return out

使用 numba:

from numba import jit, float64, int64

@jit(float64(float64[:], int64), nopython=True, parallel=True)
def fast_rosenbrock(x, N):
    out = 0.0
    for i in range(N-1):
        out += 100.0 * (x[i+1] - x[i]**2)**2 / (1 - x[i])**2
    return out

这会抛出一个错误 ZeroDivisionError: division by zero

我是不是做错了什么?

已解决错误。 numba 似乎不支持 '/' 运算符。因此,我们需要在需要的地方使用 np.divide 。以下是更新后的代码

@jit(float64(float64[:], int64), nopython=True, parallel=True)
def rosenbrock(x, N):
    out = 0.0
    for i in range(N-1):
        out += np.divide(100.0 * (x[i+1] - x[i]**2)**2, (1 - x[i])**2)
    return out

结果:

没有 Numba:78.4 ms ± 1.23 ms per loop

使用 Numba:6.59 ms ± 152 µs per loop

这几乎是计算时间的 10 倍改进。