python - 使用 Numba 进行 Numpy 优化

Question

我在球体上有两组点，在下面的代码示例中标记为“obj”和“ps”。我想识别所有距离“ps”点小于特定角度距离的“obj”点。

我对此的看法是用 3D 单位向量表示每个点，并将它们的点积与 cos(最大分离)进行比较。这可以通过 numpy 广播轻松完成，但在我的应用程序中，我有 n_obj ~ 500,000 和 n_ps ~ 50,000，因此广播的内存需求太大。下面我使用 numba 粘贴了我当前的拍摄。这可以进一步优化吗？

from numba import jit
import numpy as np
from sklearn.preprocessing import normalize

def gen_points(n):
    """
    generate random 3D unit vectors (not uniform, but irrelevant here)
    """
    vec = 2*np.random.rand(n,3)-1.
    vec_norm = normalize(vec)
    return vec_norm

#@jit(nopython=True)
@jit
def angdist_threshold_numba(vec_obj,vec_ps,cos_maxsep):
    """
    finds obj that are closer than maxsep to a ps
    """    
    nps = len(vec_ps)
    nobj = len(vec_obj)     

    #closeobj_all = []
    closeobj_all = np.empty(0)
    dotprod = np.empty(nobj)
    a = np.arange(nobj)
    for ps in range(nps):
        np.sum(vec_obj*vec_ps[ps],axis=1,out=dotprod)
        #closeobj_all.extend(a[dotprod > cos_maxsep])
        closeobj_all = np.append(closeobj_all, a[dotprod > cos_maxsep])  

    return closeobj_all


vec_obj = gen_points(50000) #in reality ~500,000
vec_ps = gen_points(5000) #in reality ~50,000
cos_maxsep = np.cos(0.003)

closeobj_all = np.unique(angdist_threshold_numba(vec_obj,vec_ps,cos_maxsep))

这是使用代码中给出的测试用例的性能:

%timeit np.unique(angdist_threshold_numba(vec_obj,vec_ps,cos_maxsep))
1 loops, best of 3: 4.53 s per loop

我尝试使用

来加快它的速度

@jit(nopython=True)

但这失败了

NotImplementedError: Failed at nopython (nopython frontend)
(<class 'numba.ir.Expr'>, build_list(items=[]))

编辑:在 numba 更新到 0.26 后，即使在 python 模式下，空列表的创建也会失败。这可以通过将其替换为 np.empty(0) 并将 .extend() 替换为 np.append() 来解决，见上文。这几乎不会改变性能。

根据 https://github.com/numba/numba/issues/858 np.empty() 现在在 nopython 模式下受支持，但我仍然无法使用 @jit(nopython = True) 运行它:

TypingError: Internal error at <numba.typeinfer.CallConstraint object at 0x7ff3114a9310>

Answer 1

与 list.append 不同，你不应该在循环中调用 numpy.append!这是因为即使追加单个元素，也需要复制整个数组。因为您只对唯一的 obj 感兴趣，所以您可以使用 bool 数组来标记到目前为止找到的匹配项。

至于 Numba，如果您写出所有循环，效果最好。例如:

@jit(nopython=True)
def numba2(vec_obj, vec_ps, cos_maxsep):
    nps = vec_ps.shape[0]
    nobj = vec_obj.shape[0]
    dim = vec_obj.shape[1]
    found = np.zeros(nobj, np.bool_)
    for i in range(nobj):
        for j in range(nps):
            cos = 0.0
            for k in range(dim):
                cos += vec_obj[i,k] * vec_ps[j,k]
            if cos > cos_maxsep:
                found[i] = True
                break
    return found.nonzero()

额外的好处是，一旦找到与当前 obj 的匹配项，我们就可以跳出 ps 数组的循环。

您可以通过专门针对 3 维空间的函数来提高速度。此外，出于某种原因，将所有数组和相关维度传递给辅助函数会导致另一个加速:

def numba3(vec_obj, vec_ps, cos_maxsep):
    nps = len(vec_ps)
    nobj = len(vec_obj)
    out = np.zeros(nobj, bool)
    numba3_helper(vec_obj, vec_ps, cos_maxsep, out, nps, nobj)
    return np.flatnonzero(out)

@jit(nopython=True)
def numba3_helper(vec_obj, vec_ps, cos_maxsep, out, nps, nobj):
    for i in range(nobj):
        for j in range(nps):
            cos = (vec_obj[i,0]*vec_ps[j,0] + 
                   vec_obj[i,1]*vec_ps[j,1] + 
                   vec_obj[i,2]*vec_ps[j,2])
            if cos > cos_maxsep:
                out[i] = True
                break
    return out

我得到 20,000 obj 和 2,000 ps 的时间:

%timeit angdist_threshold_numba(vec_obj,vec_ps,cos_maxsep)
1 loop, best of 3: 2.99 s per loop
%timeit numba2(vec_obj, vec_ps, cos_maxsep)
1 loop, best of 3: 444 ms per loop
%timeit numba3(vec_obj, vec_ps, cos_maxsep)
10 loops, best of 3: 134 ms per loop