为什么Numba没有改善这个迭代...?

2 投票
2 回答
1130 浏览
提问于 2025-04-18 01:17

我正在尝试使用Numba来加速一个计算联合发生的最小条件概率的函数。

    import numpy as np
    from numba import double
    from numba.decorators import jit, autojit

    X = np.random.random((100,2))

    def cooccurance_probability(X):
        P = X.shape[1]      
        CS = np.sum(X, axis=0)                  #Column Sums
        D = np.empty((P, P), dtype=np.float)    #Return Matrix
        for i in range(P):
            for j in range(P):
                D[i, j] = (X[:,i] * X[:,j]).sum() / max(CS[i], CS[j])
        return D 

    cooccurance_probability_numba = autojit(cooccurance_probability)

不过,我发现cooccurance_probabilitycooccurance_probability_numba的性能几乎是一样的。

%timeit cooccurance_probability(X)
1 loops, best of 3: 302 ms per loop

%timeit cooccurance_probability_numba(X)
1 loops, best of 3: 307 ms per loop

这是为什么呢?是不是因为numpy在逐个元素操作的缘故?

我参考的例子是: http://nbviewer.ipython.org/github/ellisonbg/talk-sicm2-2013/blob/master/NumbaCython.ipynb

[注意:由于问题的对称性,我可以将执行时间减半,但这不是我主要关心的]

2 个回答

1

下面是根据Josh的建议给出的一个解决方案,这个建议非常准确。不过在下面的实现中,max()函数似乎运行得很好。如果能有一个“安全”的Python / Numpy函数列表,那就太好了。

注意:我把原始矩阵的维度减少到了100 x 200。

import numpy as np
from numba import double
from numba.decorators import jit, autojit

X = np.random.random((100,200))

def cooccurance_probability_explicit(X):
    C = X.shape[0]
    P = X.shape[1]      
    # - Column Sums - #
    CS = np.zeros((P,), dtype=np.float)
    for p in range(P):
        for c in range(C):
            CS[p] += X[c,p]
    D = np.empty((P, P), dtype=np.float)    #Return Matrix
    for i in range(P):
        for j in range(P):
            # - Compute Elemental Pairwise Sums over each Product Vector - #
            pws = 0
            for c in range(C):
                pws += (X[c,i] * X[c,j])
            D[i,j] = pws / max(CS[i], CS[j])
    return D 

cooccurance_probability_explicit_numba = autojit(cooccurance_probability_explicit)

%timeit 结果:

%timeit cooccurance_probability(X)
10 loops, best of 3: 83 ms per loop


%timeit cooccurance_probability_explicit(X)
1 loops, best of 3: 2.55s per loop

%timeit cooccurance_probability_explicit_numba(X)
100 loops, best of 3: 7.72 ms per loop

有趣的是,Python直接执行的版本非常慢,因为它在检查数据类型时花费了很多时间。但通过Numba处理后,效果就很好了。(Numba的速度大约是使用Numpy的Python解决方案的11.5倍)。


更新:添加了一个Cython函数进行比较(感谢moarningsun: 带有可变大小矩阵输入的Cython函数

%load_ext cythonmagic
%%cython
import numpy as np
cimport numpy as np

def cooccurance_probability_cy(double[:,:] X):
    cdef int C, P, i, j, k
    C = X.shape[0]
    P = X.shape[1]
    cdef double pws
    cdef double [:] CS = np.sum(X, axis=0)
    cdef double [:,:] D = np.empty((P,P), dtype=np.float)

    for i in range(P):
        for j in range(P):
            pws = 0.0
            for c in range(C):
                pws += (X[c, i] * X[c, j])
            D[i,j] = pws / max(CS[i], CS[j]) 
    return D

%timeit 结果:

%timeit cooccurance_probability_cy(X)
100 loops, best of 3: 12 ms per loop
3

我猜你的问题可能是因为你在使用 sum 函数时,程序没有生成原生代码,而是停留在了对象层面。这就意味着 Numba 可能不会显著加快速度,因为它目前还不知道怎么优化或转换 sum。另外,通常来说,把向量化的操作展开成明确的循环会更好。你链接的那个 ipynb 文件里只调用了 np.sqrt,我认为这个函数是可以转换成机器代码的,而且它是对每个元素进行操作,而不是对切片进行操作。我建议你在内层循环中把 sum 展开成一个明确的循环,逐个处理元素,而不是使用切片和 sum 方法。

根据我的经验,Numba 有时能带来很大的提升,但它并不能加速所有的 Python 代码。你需要了解它的局限性,以及它能有效优化哪些部分。另外,值得注意的是,Numba 的 v0.11 版本在这方面和 0.12、0.13 版本有些不同,因为在这几个版本之间,Numba 进行了重大重构。

撰写回答