有一种信念,即numpy的功能不能通过cython或numba加速。但这并不完全正确:numpy的目标是为各种场景提供出色的性能,但这也意味着某些特殊场景下性能略低于完美。
如果手头有特定场景,您就有机会改进numpy的性能,即使这意味着重写部分numpy的功能。例如,在这种情况下,我们可以使用cython将函数加速4倍,使用numba将函数加速8倍。
让我们以您的版本作为基准(请参见答案末尾的列表):
>>>%timeit cosine(x,y) # scipy's
31.9 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>>%timeit np_cosine(x,y)
4.05 µs ± 19.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit np_cosine_fhtmitchell(x,y) # @FHTmitchell's version
4 µs ± 53.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
>>>%timeit np_cy_cosine(x,y)
2.56 µs ± 123 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
所以我看不到@FHTmitchell版本的改进,但与您的计时没有区别。
您的向量只有50个元素,因此需要大约200-300纳秒进行实际计算:其他所有内容都是调用函数的开销。减少开销的一种可能性是通过cython手动“内联”这些函数:
%%cython
from libc.math cimport sqrt
import numpy as np
cimport numpy as np
def cy_cosine(np.ndarray[np.float64_t] x, np.ndarray[np.float64_t] y):
cdef double xx=0.0
cdef double yy=0.0
cdef double xy=0.0
cdef Py_ssize_t i
for i in range(len(x)):
xx+=x[i]*x[i]
yy+=y[i]*y[i]
xy+=x[i]*y[i]
return 1.0-xy/sqrt(xx*yy)
这导致:
>>> %timeit cy_cosine(x,y)
921 ns ± 19.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
不错!我们可以通过放弃一些安全措施(运行时检查+ieee-754标准),通过进行以下更改来提高性能:
%%cython -c=-ffast-math
...
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
def cy_cosine_perf(np.ndarray[np.float64_t] x, np.ndarray[np.float64_t] y):
...
这导致:
>>> %timeit cy_cosine_perf(x,y)
828 ns ± 17.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
即另外10%,这意味着比numpy版本快了近5倍。
还有另一个工具提供类似的功能/性能-numba:
import numba as nb
import numpy as np
@nb.jit(nopython=True, fastmath=True)
def nb_cosine(x, y):
xx,yy,xy=0.0,0.0,0.0
for i in range(len(x)):
xx+=x[i]*x[i]
yy+=y[i]*y[i]
xy+=x[i]*y[i]
return 1.0-xy/np.sqrt(xx*yy)
这导致:
>>> %timeit nb_cosine(x,y)
495 ns ± 5.9 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
相比原始的numpy版本,速度提升了8倍。
Numba之所以更快,有以下一些原因:Cython在运行时处理数据的步幅,可以防止某些优化(如向量化)。而Numba似乎处理得更好。
但这里numba之所以更快完全是由于减少了开销:
%%cython -c=-ffast-math
import numpy as np
cimport numpy as np
def cy_empty(np.ndarray[np.float64_t] x, np.ndarray[np.float64_t] y):
return x[0]*y[0]
import numba as nb
import numpy as np
@nb.jit(nopython=True, fastmath=True)
def nb_empty(x, y):
return x[0]*y[0]
%timeit cy_empty(x,y)
753 ns ± 6.81 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
%timeit nb_empty(x,y)
456 ns ± 2.47 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
numba的开销几乎降低了2倍!
正如@max9111所指出的,numpy内联其他已编译函数,但也能够以很少的开销调用一些numpy函数,因此以下版本(将inner
替换为dot
):
@nb.jit(nopython=True, fastmath=True)
def np_nb_cosine(x,y):
return 1 - np.dot(x,y)/sqrt(np.dot(x,x)*np.dot(y,y))
>>> %timeit np_nb_cosine(x,y)
605 ns ± 5.9 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
只慢了约10%。
请注意,以上比较只适用于包含50个元素的向量。对于更多的元素,情况将完全不同:numpy版本使用了并行化的mkl(或类似)实现点积,将轻松击败我们简单的尝试。
这引出了一个问题:为一个特定的输入大小优化代码真的值得吗?有时答案是“是”,有时答案是“否”。
如果可能的话,我会选择
numba
+
dot
解决方案,它对于小的输入非常快,但对于大的输入也有mkl-implementation的全部功能。
还有一点小区别:第一个版本返回一个
np.float64
对象,而Cython和Numba版本返回Python浮点数。
清单:
from scipy.spatial.distance import cosine
import numpy as np
x=np.arange(50, dtype=np.float64)
y=np.arange(50,100, dtype=np.float64)
def np_cosine(x,y):
return 1 - inner(x,y)/sqrt(np.dot(x,x)*dot(y,y))
from numpy import inner, sqrt, dot
def np_cosine_fhtmitchell(x,y):
return 1 - inner(x,y)/sqrt(np.dot(x,x)*dot(y,y))
%%cython
from libc.math cimport sqrt
import numpy as np
def np_cy_cosine(x,y):
return 1 - np.inner(x,y)/sqrt(np.dot(x,x)*np.dot(y,y))