当前位置: 首页 > 工具软件 > Numba > 使用案例 >

numba安装与使用

胡霖
2023-12-01

一、numba是什么

Numba是一个针对Python的开源JIT编译器,由Anaconda公司主导开发,可以对Python原生代码进行CPU和GPU加速。Numba对NumPy数组和函数非常友好。解释器可以参考第四章内容。

CUDA C编程开发成本高,所以一般会用python先尝试加速模型训练。numba包就是python提供的GPU编程开发包。只管调用就对了。

除此之外还有很多,比方说等价于numpy的cupy,用于矩阵运算,pycuda之类的,tensorflow,pytorch等等。

参考这篇文章:numba与cupy

二、numba下载与安装

使用conda安装Numba:

$ conda install numba

或者使用pip安装:

$ pip install numba

三、numba的使用

使用方法很简单,想写好真不容易!!!

from numba import jit
import numpy as np

a = np.arange(10)
b = np.arange(1,11)

@jit
def test():
    return a+b

test()

我们只需要在原来的代码上添加一行@jit,即可将一个函数编译成机器码,其他地方都不需要更改。@符号装饰了原来的代码,所以称类似写法为装饰器。

这里装饰的函数调用的API是有限制的!!!!比如pandas是更高层次的封装,Numba其实不能理解它里面做了什么,所以无法对其加速。一些大家经常用的机器学习框架,如scikit-learn,tensorflow,pytorch等,已经做了大量的优化,不适合再使用Numba做加速。

此外,Numba不支持:try…except 异常处理,with 语句,类定义class,yield from

Numba的@jit装饰器就像自动驾驶,用户不需要关注到底是如何优化的,Numba去尝试进行优化,如果发现不支持,那么Numba会继续用Python原来的方法去执行该函数,即图 Python解释器工作原理中左侧部分。这种模式被称为object模式。

Numba真正牛逼之处在于其nopython模式。将装饰器改为@jit(nopython=True)或者@njit,Numba会假设你已经对所加速的函数非常了解,强制使用加速的方式,不会进入object模式,如编译不成功,则直接抛出异常。nopython的名字会有点歧义,我们可以理解为不使用很慢的Python,强制进入图 Python解释器工作原理中右侧部。简而言之,无脑nonpython模式就对了!!!!

注意点一

Numba库提供的是一种懒编译(Lazy Compilation)技术,即在运行过程中第一次发现代码中有@jit,才将该代码块编译。用到的时候才编译,看起来比较懒,所以叫懒编译。使用Numba时,总时间 = 编译时间 + 运行时间。相比所能节省的计算时间,编译的时间开销很小,所以物有所值。对于一个需要多次调用的Numba函数,只需要编译一次,后面再调用时就不需要编译了。

注意点二

原生Python速度慢的另一个重要原因是变量类型不确定。声明一个变量的语法很简单,如a = 1,但没有指定a到底是一个整数和一个浮点小数。Python解释器要进行大量的类型推断,会非常耗时。同样,引入Numba后,Numba也要推断输入输出的类型,才能转化为机器码。针对这个问题,Numba给出了名为Eager Compilation的优化方式。

from numba import jit, int32

@jit("int32(int32, int32)", nopython=True)
def f2(x, y):
    return x + y

@jit(int32(int32, int32))告知Numba你的函数在使用什么样的输入和输出,括号内是输入,括号左侧是输出。这样不会加快执行速度,但是会加快编译速度,可以更快将函数编译到机器码上。

具体用法参考皮皮鲁的科技星球

调用cuda,不需要选择njit修饰器,用cuda.jit即可

from numba import cuda
@cuda.jit
def gpu_add(a, b, result, N):
    idxWithinGrid = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x 
    gridStride = cuda.gridDim.x * cuda.blockDim.x
    for i in range(idxWithinGrid, N, gridStride):
        result[i] = a[i] + b[i]

四、numba的原理

python慢的原因主要是因为解释器。解决办法有俩个,一种解决办法是使用C/C++语言重写Python函数,但是这要求程序员对C/C++语言熟悉,且调试速度慢,不适合绝大多数Python程序员。另外一种非常方便快捷的解决办法就是使用Just-In-Time(JIT)技术。

Just-In-Time(JIT)技术为解释语言提供了一种优化,它能克服上述效率问题,极大提升代码执行速度,同时保留Python语言的易用性。使用JIT技术时,JIT编译器将Python源代码编译成机器直接可以执行的机器语言,并可以直接在CPU等硬件上运行。这样就跳过了原来的虚拟机,执行速度几乎与用C语言编程速度并无二致。

 类似资料: