考虑一个维度为 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 维数的中间 3D 数组时有一些开销,所以性能更差。
我试过 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 相同的性能。我想知道如果以正确的方式完成,是否有可能实现一些显着的性能改进?或者是否有其他方法可以加快速度(不涉及并行计算)?
最佳答案
Cython 的关键是尽可能避免使用 Python 对象和函数调用,包括对 numpy 数组的矢量化操作。这通常意味着手动写出所有循环并一次对单个数组元素进行操作。
有一个 very useful tutorial here涵盖了将 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
关于python - 使用 Numpy 和 Cython 加速距离矩阵计算,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/25213603/