Numpy向量化与加速

Question

Numpy向量化与加速

4

我发现了一个小的代码片段，它曾经使用了一个双重循环，我成功地用向量化将其变成了单个循环。这样做导致时间显著提高，因此我想知道是否可能通过向量化来消除第二个循环，以及是否会改善性能。

import numpy as np
from timeit import default_timer as timer
nlin, npix = 478, 480
bb = np.random.rand(nlin,npix)
slope = -8
fac = 4
offset= 0
barray = np.zeros([2,2259]);

timex = timer()
for y in range(nlin):
    for x in range(npix):
        ling=(np.ceil((x-y/slope)*fac)+1-offset).astype(np.int);
        barray[0,ling] +=1;
        barray[1,ling] +=bb[y,x];
newVar = np.copy(barray)
print(timer() - timex)

所以可以通过创建以下矩阵来将ling取出循环。

lingMat = (np.ceil((np.vstack(npixrange)-nlinrange/slope)*fac)+1-offset).astype(np.int);

这可以满足 lingMat[x,y] = "在 x 和 y 的循环中的ling"。这为向量化迈出了第一步。

- Sadikov

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Yacola · Accepted Answer

在向量化方面，你可以使用基于np.add.at的内容: np.add.at。

def yaco_addat(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    nlin_range = np.arange(nlin)
    npix_range = np.arange(npix)
    ling_mat = (np.ceil((npix_range-nlin_range[:,None]/slope)*fac)+1-offset).astype(np.int)  
    np.add.at(barray[0,:],ling_mat,1)
    np.add.at(barray[1,:],ling_mat,bb) 
    return barray

不过，我建议你直接使用numba来进行优化，使用带有选项nopython=True的@jit修饰器，这将使你获得：

import numpy as np
from numba import jit

nlin, npix = 478, 480
bb = np.random.rand(nlin,npix)
slope = -8
fac = 4
offset= 0

def yaco_plain(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    for y in range(nlin):
        for x in range(npix):
            ling=(np.ceil((x-y/slope)*fac)+1-offset).astype(np.int)
            barray[0,ling] += 1
            barray[1,ling] += bb[y,x]
    return barray

@jit(nopython=True)
def yaco_numba(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    for y in range(nlin):
        for x in range(npix):
            ling = int((np.ceil((x-y/slope)*fac)+1-offset))
            barray[0,ling] += 1
            barray[1,ling] += bb[y,x]    
    return barray

让我们检查输出

np.allclose(yaco_plain(bb,slope,fac,offset),yaco_addat(bb,slope,fac,offset))
>>> True
np.allclose(yaco_plain(bb,slope,fac,offset),yaco_jit(bb,slope,fac,offset))
>>> True

现在是时候进行计时了。

%timeit yaco_plain(bb,slope,fac,offset)
>>> 648 ms ± 4.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit yaco_addat(bb,slope,fac,offset)
>>> 27.2 ms ± 92.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit yaco_jit(bb,slope,fac,offset)
>>> 505 µs ± 995 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

导致函数得到优化，比起最初的两个循环版本，速度要快得多，比起np.add.at，快了53x倍。希望这对你有所帮助。