(2026.03.31 Tues)
Numba是Python的即時(shí)編譯器(just-in-time JIT compiler)，特別適用于使用了Numpy array的函數(shù)和循環(huán)，即computation bound的任務(wù)。

使用Numba最常見的方式是其內(nèi)部的一系列裝飾器，使被裝飾的函數(shù)由Numba編譯。Numba裝飾器會(huì)讀取被裝飾函數(shù)的python字節(jié)碼(bytecode)，并將其與函數(shù)輸入變量類型相連(?)。經(jīng)過分析和優(yōu)化后，由LLVM編譯器庫(complier library)生成函數(shù)的機(jī)器碼，繞過python解釋器，并以原生機(jī)器速度運(yùn)行。之后每次調(diào)用該函數(shù)時(shí)都使用該機(jī)器碼。

案例1：計(jì)算圓周率

用Monte Carlo法計(jì)算圓周率pi，在計(jì)算過程中加入@numba.jit裝飾器，會(huì)提高運(yùn)行效率。

import functools, random
from numba import jit

# 設(shè)定計(jì)時(shí)器裝飾器
def timer(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        starting_time = time.perf_counter()
        result = func(*args, **kwargs)
        ending_time = time.perf_counter()
        print('time consumption in seconds: ', ending_time - starting_time)
        return result
    return wrapper

@timer
@jit(nopython=True)
def mc_pi_jit(nsamples):
    acc = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

運(yùn)行

>> mc_pi_jit(10000000)
time consumption in seconds:  0.31119772199963336

調(diào)用@jit裝飾器時(shí)，指定了傳遞的參數(shù)nopython=True(默認(rèn)設(shè)置)，以保證被裝飾函數(shù)在編譯時(shí)完全沒需python解釋器的接入。這個(gè)設(shè)定將極大提高代碼運(yùn)行效率。

作為對(duì)比，運(yùn)行該代碼并且不加入@jit裝飾器的運(yùn)行時(shí)長(zhǎng)如下

@timer
def mc_pi_no_jit(nsamples):
    acc = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

>> mc_pi_no_jit(10000000)
time consumption in seconds:  3.635518551000132
3.1409716

時(shí)長(zhǎng)對(duì)比：使用@jit裝飾器，0.31s；不使用@jit裝飾器，3.64s。效率提升達(dá)到10倍。

案例2：計(jì)算矩陣乘法

使用numba中的cuda，將矩陣乘法部署到GPU上。

# cuda_test.py
import numpy as np
import time
from tqdm import trange
from numba import cuda
cuda.select_device(1)

@cuda.jit
def CudaSquare(x):
    i, j = cuda.grid(2)
    x[i][j] *= x[i][j]

if __name__ == '__main__':
    numpy_time = 0
    numba_time = 0
    test_length = 1000
    for i in trange(test_length):
        np.random.seed(i)
        array_length = 2**12
        random_array = np.random.rand(array_length, array_length)
        random_array_cuda = cuda.to_device(random_array)
        time0 = time.time()
        square_array = np.square(random_array)
        time1 = time.time()
        CudaSquare[(array_length,array_length),(1,1)](random_array_cuda)
        time2 = time.time()
        numpy_time += time1-time0
        numba_time += time2-time1
    print ('The time cost of numpy is {}s for {} loops'.format(numpy_time, test_length))
    print ('The time cost of numba is {}s for {} loops'.format(numba_time, test_length))

運(yùn)行，對(duì)比部署與否的效率，運(yùn)行時(shí)間相差達(dá)到10倍以上。

>> python cuda_test.py
The time cost of numpy is 4.878739595413208s for 100 loops
The time cost of numba is 0.3255774974822998s for 100 loops

Reference

• Numba官網(wǎng)
• cnblog, DECHIN, 超過Numpy的速度有多難？試試Numba的GPU加速