使用numba对python数组复杂运算加速1000倍以上

python

发布日期: 2022-11-12

numba介绍

Numba 是适用于 Python 的即时编译器（jit），最适用于使用 NumPy 数组和函数以及循环的代码。使用 Numba 最常见的方法是通过它的装饰器集合，通过这些装饰器来标记函数使用 Numba 编译。当调用 Numba 修饰函数时，它会“即时”编译为机器代码以供执行，最后使用 LLVM 编译器库生成函数的机器码，每次调用函数时都会使用此编译版本。

使用案例

在使用numba进行加速时，首先需要在待加速函数前添加@git装饰器，但这并不足够。因为numba中支持的数据类别和运算类型，和numpy不完全一致，具体可参考numba supported types。

因此，一般情况下需要对原来的函数代码进行修改。涉及的修改情况，比如list类型，比如numpy运算np.zeros()等，需要将这些修改为numba中支持的类型。只有添加了jit编译装饰器，并且将代码完全修改为numba所支持的类型和运算，那么通过numba加速才能实现。

下面是列的一个numba加速的参考示例，对于数据类型和运算均有进行相应修改匹配，使其符合numba。

from numba.typed import List


@jit(nopython=True, parallel=True)
def my_fun(ary_data, nums = 500, ab_ratio=0.01):
    rows = ary_data.shape[0]
    stmp = np.zeros((rows, 1), dtype=np.int64)
    for i in range(nums):
        r_num = int(rows * ab_ratio)
        ind = np.arange(rows)
        index = ind[0:r_num]  
        val = ary_data[index,:]               
        val_mean = np.mean(val) 
        val_std = np.std(val)                
        tt = []
        for j in range(rows):
            rr = (ary_data[j,:] - val_mean) / val_std
            tmp = np.sqrt(rr)
            tt.append(tmp)
    return np.array(tt)


@jit(nopython=True, parallel=True)
def use_fun(data_src, class_num, nums = 500, ab_ratio=0.01):
    inds_final = List()
    for cc in prange(class_num):
        inds_level = np.where(data_src[:,-1]==cc)[0]
        data_level = data_src[inds_level]
        selectIdx = my_fun(data_level, nums, ab_ratio)

        resIdx = [inds_level[v] for v in selectIdx]
        inds_final.append(resIdx)
   
    return inds_final

df_data_use = df[names+[target]].values
inds_select = use_fun(df_data_use, labes_nums, 100)