我的 numba 代码能比 numpy 快吗

Question

我是 Numba 的新手，我正在尝试加速一些已证明对 numpy 来说太笨重的计算。我在下面给出的示例比较了一个包含我的计算子集的函数，该函数使用函数的 vectorized/numpy 和 numba 版本，后者也通过注释掉 @autojit 装饰器被测试为纯 python .

我发现 numba 和 numpy 版本相对于纯 python 提供了相似的加速，两者都是大约 10 倍的速度提升。 numpy 版本实际上比我的 numba 函数稍快，但由于此计算的 4D 性质，当 numpy 函数中的数组的大小比这个玩具示例大得多时，我很快运行内存不足。

这种提速很好，但我经常在网络上看到从纯 python 迁移到 numba 时提速超过 100 倍。

我想知道在 nopython 模式下移动到 numba 时是否有普遍预期的速度提升。我还想知道我的 numba 化函数中是否有任何组件会限制进一步的速度提升。

import numpy as np                                                                    
from timeit import default_timer as timer                                             
from numba import  autojit                                                            
import math                                                                           


def vecRadCalcs(slope,  skyz, solz, skya, sola):                                      

    nloc = len(slope)                                                                 
    ntime =  len(solz)                                                                
    [lenz, lena] = skyz.shape                                                         
    asolz = np.tile(np.reshape(solz,[ntime,1,1,1]),[1,nloc,lenz,lena])                
    asola = np.tile(np.reshape(sola,[ntime,1,1,1]),[1,nloc,lenz,lena])                
    askyz = np.tile(np.reshape(skyz,[1,1,lenz,lena]),[ntime,nloc,1,1])                
    askya = np.tile(np.reshape(skya,[1,1,lenz,lena]),[ntime,nloc,1,1])                
    phi1 = np.cos(asolz)*np.cos(askyz)                                                
    phi2 = np.sin(asolz)*np.sin(askyz)*np.cos(askya- asola)                           
    phi12 = phi1 + phi2                                                               
    phi12[phi12> 1.0] = 1.0                                                           
    phi = np.arccos(phi12)                                                            

    return(phi)                                                                       


@autojit                                                                              
def RadCalcs(slope,  skyz, solz, skya, sola, phi):                                    

    nloc = len(slope)                                                                 
    ntime =  len(solz)                                                                
    pop = 0.0                                                                         
    [lenz, lena] = skyz.shape                                                         
    for iiT in range(ntime):                                                          
        asolz = solz[iiT]                                                             
        asola = sola[iiT]                                                             
        for iL in range(nloc):                                                        
            for iz in range(lenz):                                                    
                for ia in range(lena):                                                
                    askyz = skyz[iz,ia]                                               
                    askya = skya[iz,ia]                                               
                    phi1 = math.cos(asolz)*math.cos(askyz)                            
                    phi2 = math.sin(asolz)*math.sin(askyz)*math.cos(askya- asola)     
                    phi12 = phi1 + phi2                                               
                    if phi12 > 1.0:                                                   
                        phi12 = 1.0                                                   
                    phi[iz,ia] = math.acos(phi12)                                     
                    pop = pop + 1                                                     

    return(pop)                                                                       


zenith_cells = 90                                                                     
azim_cells = 360                                                                      
nloc = 10        # nominallly ~ 700                                                   
ntim = 10        # nominallly ~ 200000                                                

slope = np.random.rand(nloc) * 10.0                                                   
solz = np.random.rand(ntim) *np.pi/2.0                                                
sola = np.random.rand(ntim) * 1.0*np.pi                                               

base = np.ones([zenith_cells,azim_cells])                                             
skya =  np.deg2rad(np.cumsum(base,axis=1))                                            
skyz =  np.deg2rad(np.cumsum(base,axis=0)*90/zenith_cells)                            

phi = np.zeros(skyz.shape)                                                            

start = timer()                                                                       
outcalc = RadCalcs(slope,  skyz, solz, skya, sola, phi)                               
stop = timer()                                                                        
outcalc2 = vecRadCalcs(slope,  skyz, solz, skya, sola)                                
stopvec = timer()                                                                     

print(outcalc)                                                                        
print(stop-start)                                                                     
print(stopvec-stop)

Answer 1

在我的机器上运行ning numba 0.31.0，Numba 版本比矢量化解决方案快 2 倍。对 numba 函数进行计时时，您需要多次运行函数，因为第一次看到 jitting 代码的时间 + 运行时间。随后的运行s 将不包括 jit 函数时间的开销，因为 Numba 将 jitted 代码缓存在内存中。

另外，请注意，您的函数计算的不是同一件事——您要小心，您在比较相同的事情时使用类似 np.allclose 的结果。

我的 numba 代码能比 numpy 快吗

Can my numba code be faster than numpy

python

numpy

numba