如何使用Cython重写几个NumPy代码？

Question

如何使用Cython重写几个NumPy代码？

pythonnumpycythonnumba

3

我只是想加速使用numpy编写的数值算法。关键部分是计算对数似然函数（两个截断正态CDF之间的差异）。我的函数非常慢（每次循环31.9毫秒），而且我需要在每次迭代中运行它2000次。

我尝试使用scipy的“norm.cdf”函数代替“ecfc”。但它更慢。我还尝试过Numba包中的“@jit”。但它比原始代码还要慢。

我认为也许我需要使用Cython。但我对C几乎一无所知。我试图从Cython for numpy users网页学习Cython，但对我来说太难了。

有人能帮我用Cython重新编写代码吗？或者建议我如何编写更快的代码？

import numpy as np 
from scipy.special import erfc
# The bloody function for calculating the difference between two truncated normal CDFs
def my_loglikelihood2(x,b,c,z):
    log_likelihood=np.zeros(np.shape(z)[0])
    log_likelihood[x==1]=np.log(0.5*erfc(-(c[1]-np.dot(z[x==1,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[0]-np.dot(z[x==1],b)) / np.sqrt(2.)))
    log_likelihood[x==2]=np.log(0.5*erfc(-(c[2]-np.dot(z[x==2,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[1]-np.dot(z[x==2],b)) / np.sqrt(2.)))
    log_likelihood[x==3]=np.log(0.5*erfc(-(c[3]-np.dot(z[x==3,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[2]-np.dot(z[x==3],b)) / np.sqrt(2.)))
    return log_likelihood

# generate random values
x=np.random.randint(low=1, high=4, size=50000)
b=np.random.normal(0,1,70)
c=np.array([-999,-1,1,999],dtype='f')  
z=np.random.multivariate_normal(np.zeros(70), np.eye(70), 50000) 

%timeit my_loglikelihood2(x,b,c,z)

# 10 loops, best of 3: 31.9 ms per loop   :(

更新 1 - 基于 @jackvdp 的建议。已经加速了 4.5 倍。但我仍在寻找更快的代码：

def up_cutoff(x,c):
    x[x==1]=c[1]
    x[x==2]=c[2]
    x[x==3]=c[3] 
    return x
def low_cutoff(x,c):
    x[x==1]=c[0]
    x[x==2]=c[1]
    x[x==3]=c[2] 
def my_loglikelihood2(x,b,low_c,up_c,z):
    up_c=up_cutoff(x,c)
    low_c=low_cutoff(x,c)
    return np.log(0.5*erfc(-(up_c-np.dot(z,b)) / np.sqrt(2.)) - 0.5*erfc(-(low_c-np.dot(z,b)) / np.sqrt(2.))) 

%timeit my_loglikelihood2(x,b,low_c,up_c,z)
100 loops, best of 3: 6.58 ms per loop

更新2 - 基于@DSM的建议。将np.dot(z,b)替换为zdotb = z.dot(b)。提高了1.5ms

def my_loglikelihood2(x,b,low_c,up_c,z):
    up_c=up_cutoff(x,c)
    low_c=low_cutoff(x,c)
    zdotb = z.dot(b)
    return np.log(0.5*erfc(-(up_c-zdotb) / np.sqrt(2.)) - 0.5*erfc(-(low_c-zdotb) / np.sqrt(2.))) 

%timeit my_loglikelihood2(x,b,low_c,up_c,z)
100 loops, best of 3: 5.02 ms per loop

- sc10mmj

1

看起来常见子表达式消除和高级索引（c[x]和c[x-1]）可以大大加快这个程序的速度，几乎接近使用Cython获得的速度。 - user2357112

1

可能是这样的zdotb = z.dot(b); return np.log(0.5*erfc(-(c[x]-zdotb) / np.sqrt(2.)) - 0.5*erfc(-(c[x-1]-zdotb) / np.sqrt(2.)))。但这可能不准确 - 我现在没有测试和调试NumPy代码的设置。 - user2357112

@jakevdp 谢谢。我已经更新了代码，使用一行代码计算log_likelihood而不是3行。现在速度快了4.5倍。但我仍然想进一步加速。 - sc10mmj

正如@user2357112所指出的那样，您没有消除常见的子表达式。即使在您的新代码中，您也因某种原因计算了z.dot(b)两次。 - DSM

@DSM 谢谢。我已经更改了这个。它已经加速了1.5毫秒。对我来说是一个很大的改进。谢谢。 - sc10mmj

显示剩余5条评论

1个回答

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

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

- hpaulj · Accepted Answer

如果你的Python代码因为循环而变慢，那么将其转换为Cython可以看到巨大的改进。但是你的样本只调用现有的numpy/scipy函数几次。

主要是对np.log, erfc, np.dot, np.sqrt的调用。我不确定erfc，但其他函数已经使用编译好的代码。Cython不会触及这些函数。

我们可以检查一下erfc。

但最好的方法是使用更大的数组来调用此代码。