使用Numpy和Cython加速距离矩阵计算

Question

使用Numpy和Cython加速距离矩阵计算

pythonperformanceoptimizationnumpycython

9

考虑一个维度为NxM的numpy数组A。目标是计算欧几里得距离矩阵D，其中每个元素D[i,j]是第i行和第j行之间的欧几里得距离。最快的方法是什么？这并不是我需要解决的问题，但这是一个很好的例子，说明我想要做的事情（通常，可以使用其他距离度量）。

到目前为止，这是我能想到的最快方法：

n = A.shape[0]
D = np.empty((n,n))
for i in range(n):
    D[i] = np.sqrt(np.square(A-A[i]).sum(1))

但这是最快的方法吗？我主要关心for循环。我们能用Cython之类的工具击败它吗？

为了避免循环，我尝试使用广播，并进行以下操作：

D = np.sqrt(np.square(A[np.newaxis,:,:]-A[:,np.newaxis,:]).sum(2))

但是事实证明这是一个不好的想法，因为构建一个中间的三维数组NxNxM会有一些开销，所以性能会更差。

我尝试了Cython。但我在Cython方面是新手，所以我不知道我的尝试有多好：

def dist(np.ndarray[np.int32_t, ndim=2] A):
    cdef int n = A.shape[0]    
    cdef np.ndarray[np.float64_t, ndim=2] dm = np.empty((n,n), dtype=np.float64)      
    cdef int i = 0    
    for i in range(n):  
        dm[i] = np.sqrt(np.square(A-A[i]).sum(1)).astype(np.float64)              
    return dm

上面的代码比Python的for循环慢了一些。我不太了解Cython，但我认为我至少可以达到与for循环+numpy相同的性能。我想知道在正确的方式下是否有可能实现一些显著的性能提升？或者是否有其他方法可以加速（不涉及并行计算）？

- ojy

1

此外，为了实现这一点，使用Cython编写ufunc，然后将其应用于'A'可能会更容易，而不是将整个循环放在Cython中。这样做错误的可能性较小，如果没有其他问题的话... - abarnert

6

有一个专门执行此任务的SciPy方法（http://docs.scipy.org/doc/scipy-0.13.0/reference/generated/scipy.spatial.distance.pdist.html），因此这可能是一个非常快速的选择。 - user2357112

@abarnert，N和M可能相当大，比如N约为20k，M约为1k。 - ojy

2

关于Cython，如果您正在使用它，您可能希望自己进行数学计算，而不是调用NumPy例程。当您已经编写可以编译为C的代码时，NumPy向量化并没有太大帮助。 - user2357112

@user2357112，这是一个有用的提示，谢谢。 - ojy

显示剩余4条评论

1个回答

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

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

- ali_m · Accepted Answer

使用Cython的关键是尽可能避免使用Python对象和函数调用，包括numpy数组上的向量化操作。这通常意味着手动编写所有循环并逐个操作单个数组元素。

这里有一个非常实用的教程，介绍了将numpy代码转换为Cython并进行优化的过程。

以下是更优化的Cython版本的距离函数：

import numpy as np
cimport numpy as np
cimport cython

# don't use np.sqrt - the sqrt function from the C standard library is much
# faster
from libc.math cimport sqrt

# disable checks that ensure that array indices don't go out of bounds. this is
# faster, but you'll get a segfault if you mess up your indexing.
@cython.boundscheck(False)
# this disables 'wraparound' indexing from the end of the array using negative
# indices.
@cython.wraparound(False)
def dist(double [:, :] A):

    # declare C types for as many of our variables as possible. note that we
    # don't necessarily need to assign a value to them at declaration time.
    cdef:
        # Py_ssize_t is just a special platform-specific type for indices
        Py_ssize_t nrow = A.shape[0]
        Py_ssize_t ncol = A.shape[1]
        Py_ssize_t ii, jj, kk

        # this line is particularly expensive, since creating a numpy array
        # involves unavoidable Python API overhead
        np.ndarray[np.float64_t, ndim=2] D = np.zeros((nrow, nrow), np.double)

        double tmpss, diff

    # another advantage of using Cython rather than broadcasting is that we can
    # exploit the symmetry of D by only looping over its upper triangle
    for ii in range(nrow):
        for jj in range(ii + 1, nrow):
            # we use tmpss to accumulate the SSD over each pair of rows
            tmpss = 0
            for kk in range(ncol):
                diff = A[ii, kk] - A[jj, kk]
                tmpss += diff * diff
            tmpss = sqrt(tmpss)
            D[ii, jj] = tmpss
            D[jj, ii] = tmpss  # because D is symmetric

    return D

我把它保存在一个名为fastdist.pyx的文件中。我们可以使用pyximport来简化构建过程:

import pyximport
pyximport.install()
import fastdist
import numpy as np

A = np.random.randn(100, 200)

D1 = np.sqrt(np.square(A[np.newaxis,:,:]-A[:,np.newaxis,:]).sum(2))
D2 = fastdist.dist(A)

print np.allclose(D1, D2)
# True

至少它能够工作。让我们使用%timeit魔法进行一些基准测试：

%timeit np.sqrt(np.square(A[np.newaxis,:,:]-A[:,np.newaxis,:]).sum(2))
# 100 loops, best of 3: 10.6 ms per loop

%timeit fastdist.dist(A)
# 100 loops, best of 3: 1.21 ms per loop

一个 ~9倍的加速是不错，但并不能改变游戏规则。正如你所说，广播方法的主要问题在于构建中间数组需要大量内存。

A2 = np.random.randn(1000, 2000)
%timeit fastdist.dist(A2)
# 1 loops, best of 3: 1.36 s per loop

我不建议使用广播来尝试这个...

我们可以通过使用prange函数将其并行化到最外层循环。

from cython.parallel cimport prange

...

for ii in prange(nrow, nogil=True, schedule='guided'):
...

为了编译并行版本，您需要告诉编译器启用OpenMP。我还没有弄清楚如何使用pyximport来实现这一点，但如果您使用的是gcc，则可以手动编译，如下所示：

$ cython fastdist.pyx
$ gcc -shared -pthread -fPIC -fwrapv -fopenmp -O3 \
   -Wall -fno-strict-aliasing  -I/usr/include/python2.7 -o fastdist.so fastdist.c

使用8个线程并行处理：

%timeit D2 = fastdist.dist_parallel(A2)
1 loops, best of 3: 509 ms per loop