Numba 优化函数没有速度增益

Question

我编写了一些代码来计算不同卡车对 300 英尺长的桥梁施加的弯矩。卡车数据包含在两个列表中：ax_list和sp_list，分别是轴重和轴距。

代码没什么大不了的，但是，这需要对数百万种不同的卡车类型重复，我正在尝试优化我的代码，当涉及实际数据集时，这需要很长时间。

我尝试使用 Numba 来查看是否可以提高速度，但它并没有改变执行时间，无论我是否为每个函数添加 Numba @jit 装饰器。我在这里做错了什么？欢迎任何帮助！我还在下面包含了为 1000 条记录生成代表性伪数据的代码：

import random
from numba import jit
import numpy as np
from __future__ import division

#Generate Random Data Set

ax_list=[]
sp_list=[]

for i in xrange(1000):
    n = random.randint(3,10)
    ax = []
    sp = [0]
    for i in xrange(n):
        a = round(random.uniform(8,32),1)
        ax.append(a)
    for i in xrange(n-1):
        s = round(random.uniform(4,30), 1)
        sp.append(s)
    ax_list.append(ax)
    sp_list.append(sp)

#Input Parameters
L=300
step_size=4
cstep_size=4
moment_list=[]

@jit
#Simple moment function
def Moment(x):
    if x<L/2.0:
        return 0.5*x
    else:
        return 0.5*(L-x)

#Attempt to vectorize the Moment function, hoping for speed gains
vectMoment = np.vectorize(Moment,otypes=[np.float],cache=False)

@jit
#Truck movement function that uses the vectorized Moment function above
def SimpleSpanMoment(axles, spacings, step_size):
    travel = L + sum(spacings)
    spacings=list(spacings)
    maxmoment = 0
    axle_coords =(0-np.cumsum(spacings))
    while np.min(axle_coords) < L:
        axle_coords = axle_coords + step_size
        moment_inf = np.where((axle_coords >= 0) & (axle_coords <=L), vectMoment(axle_coords), 0)
        moment = sum(moment_inf * axles)
        if maxmoment < moment:
            maxmoment = moment
    return maxmoment

然后运行循环1000次：

%%timeit
for i in xrange(len(ax_list)):
    moment_list.append(np.around(SimpleSpanMoment(ax_list[i], sp_list[i], step_size),1))

产量：

1 loop, best of 3: 2 s per loop

我也试过在 jit 装饰器中声明类型，但结果仍然没有变化。

@jit('f8(f8)')和@jit('f8(f8[:],f8[:],f8)')，分别为两个函数。

Answer 1

基本问题是，当您使用 nb.jit 并且遇到无法编译为本机代码的内容时，它会改用 object mode，这可能会很慢。如果将 nopython=True 指定为大多数 numba decorators/functions 的参数，则可以看到这一点。然后，如果 numba 无法显式键入变量或不知道如何转换函数，您将收到错误消息。这是一个版本，我相信它会产生与您的原始功能相同的结果。在我的机器上，您的代码需要大约 2.7 秒才能运行。完全处于 nopython 模式下运行s 下面的优化需要大约 50 毫秒（~50 倍加速）：

@nb.vectorize(nopython=True)
#Simple moment function
def vectMoment2(x):
    if x<L/2.0:
        return 0.5*x
    else:
        return 0.5*(L-x)

@nb.jit(nopython=True)
#Truck movement function that uses the vectorized Moment function above
def SimpleSpanMoment2(axles, spacings, step_size):
    travel = L + np.sum(spacings)
    maxmoment = 0
    axle_coords = -np.cumsum(spacings)

    moment_inf = np.empty_like(axles)
    while np.min(axle_coords) < L:
        axle_coords = axle_coords + step_size
        y = vectMoment2(axle_coords)

        for k in range(y.shape[0]):
            if axle_coords[k] >=0 and axle_coords[k] <= L:
                moment_inf[k] = y[k]
            else:
                moment_inf[k] = 0.0

        moment = np.sum(moment_inf * axles)
        if maxmoment < moment:
            maxmoment = moment
    return maxmoment

然后通过以下方式计时：

%%timeit
for i in xrange(len(ax_list)):
    moment_list2.append(np.around(SimpleSpanMoment2(np.array(ax_list[i]), np.array(sp_list[i]), step_size),1))

查看文档，了解 Numba 在 nopython 模式下支持的内容：

请注意，您可以在 nopython 模式下在 numba 函数内部使用 np.where，但第三个参数必须是数组（例如 np.zeros_like(moment_inf)）而不是整数。我发现它比我在上面明确循环数组的函数编码方式慢了大约 2 倍。

Numba 优化函数没有速度增益

No speed gain from Numba optimized functions

python

optimization

performance

numpy

numba