与cython的Python快速余弦距离

Question

我想尽可能加快余弦距离的计算scipy.spatial.distance.cosine,所以我尝试使用numpy

def alt_cosine(x,y):
    return 1 - np.inner(x,y)/np.sqrt(np.dot(x,x)*np.dot(y,y))

我试过cython

from libc.math cimport sqrt
def alt_cosine_2(x,y):
    return 1 - np.inner(x,y)/sqrt(np.dot(x,x)*np.dot(y,y))

并逐渐得到改进(测试长度为50的numpy数组)

>>> cosine() # ... make some timings
5.27526156300155e-05 # mean calculation time for one loop

>>> alt_cosine() 
9.913400815003115e-06

>>> alt_cosine_2()
7.0269494536660205e-06

最快的方法是什么？不幸的是,我无法指定变量类型alt_cosine_2,我将使用带有类型的numpy数组的此函数np.float32

Answer 1

有一种观点认为,在cython或numba的帮助下,numpy的功能无法加速.但这并不完全正确:numpy的目标是为各种场景提供出色的性能,但这也意味着对于特殊场景而言,性能稍差.

有了特定的场景,你就有机会改善numpy的性能,即使它意味着重写一些numpy的功能.例如,在这种情况下,我们可以使用numba使用cython和factor 8将因子加速4倍.

让我们从您的版本开始作为基线(请参阅答案末尾的列表):

>>>%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)  # your numpy-version
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 ns:其他一切都是调用函数的开销.减少开销的一种可能性是在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内联其他jitted函数,但它能够以非常小的开销调用一些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-version使用dot-product的并行化mkl(或类似)实现,并且将轻松击败我们的简单尝试.

这引出了一个问题:是否真的值得为特殊大小的输入优化代码？有时答案是肯定的,有时答案是"不".

如果有可能,我会选择numba+ dot解决方案,这对于小输入非常快,但对于更大的输入也具有mkl实现的全部功能.

还有一点不同:第一个版本返回np.float64-object,cython和numba版本返回Python-float.

人数:

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))