Python 与 cython 的快速余弦距离

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

加速这种代码的懒惰方法：

使用 numexpr Python 模块
使用 numba Python 模块
使用 SciPy 等价的 NumPy 函数

不幸的是，none 这些技巧对您有用，因为：

dot 和 inner 未在 numexpr
numba（像 Cython）不会加速调用 NumPy 的函数
dot 和 inner 在 scipy 中没有不同的实现（它们甚至在命名空间中不可用）。

也许你最好的选择是尝试在不同的底层 LA 库（例如 LAPACK、BLAS、OpenBLAS 等）和编译选项（例如多线程等）下编译 numpy 以查看哪种组合对您的用例最有效。

祝你好运！

Answer 2

人们相信，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)  # 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 的帮助下每手 "inline" 这些函数：

%%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 倍。

还有另一个工具提供类似的 functionality/performance - 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 版本使用点积的并行化 mkl（或类似）实现，将轻松击败我们的简单尝试。

这引出了一个问题：是否真的值得针对特殊大小的输入优化代码？有时答案是 "yes" 有时答案是 "no".

如果可能的话，我会得到 numba + dot 解决方案，它对于小输入非常快，但对于更大的输入也具有 mkl 实现的全部功能。

也有细微差别：第一个版本 return 一个 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))

Python 与 cython 的快速余弦距离

Python fast cosine distance with cython

python

trigonometry

numpy

cython