Python中数学函数的优化与加速

Question

Python中数学函数的优化与加速

10

这个数学函数的目的是使用二面角计算两个（或多个）蛋白质结构之间的距离：

它在结构生物学中非常有用。我已经用numpy编写了这个函数，但目标是拥有更快的实现。作为计算时间参考，我使用了scikit-learn包中可用的欧几里德距离函数。

这是我目前的代码：

import numpy as np
import numexpr as ne
from sklearn.metrics.pairwise import euclidean_distances

# We have 10000 structures with 100 dihedral angles
n = 10000
m = 100

# Generate some random data
c = np.random.rand(n,m)
# Generate random int number
x = np.random.randint(c.shape[0])

print c.shape, x

# First version with numpy of the dihedral_distances function
def dihedral_distances(a, b):
    l = 1./a.shape[0]
    return np.sqrt(l* np.sum((0.5)*(1. - np.cos(a-b)), axis=1))

# Accelerated version with numexpr
def dihedral_distances_ne(a, b):
    l = 1./a.shape[0]
    tmp = ne.evaluate('sum((0.5)*(1. - cos(a-b)), axis=1)')
    return ne.evaluate('sqrt(l* tmp)')

# The function of reference I try to be close as possible 
# in term of computation time
%timeit euclidean_distances(c[x,:], c)[0]
1000 loops, best of 3: 1.07 ms per loop

# Computation time of the first version of the dihedral_distances function
# We choose randomly 1 structure among the 10000 structures.
# And we compute the dihedral distance between this one and the others
%timeit dihedral_distances(c[x,:], c)
10 loops, best of 3: 21.5 ms per loop

# Computation time of the accelerated function with numexpr
%timeit dihedral_distances_ne(c[x,:], c)
100 loops, best of 3: 9.44 ms per loop

9.44毫秒非常快，但如果需要运行一百万次，它就非常慢。现在的问题是，下一步该怎么做？是Cython？还是PyOpenCL？我有一些PyOpenCL的经验，但我从未编写过像这样复杂的代码。我不知道是否可以像使用numpy那样在GPU上一次计算二面角距离，并且如何继续操作。

谢谢你的帮助！

编辑：谢谢大家！我正在全力解决方案，一旦完成我会把代码放在这里。

CYTHON版本：

%load_ext cython
import numpy as np

np.random.seed(1234)

n = 10000
m = 100

c = np.random.rand(n,m)
x = np.random.randint(c.shape[0])

print c.shape, x

%%cython --compile-args=-fopenmp --link-args=-fopenmp --force

import numpy as np
cimport numpy as np
from libc.math cimport sqrt, cos
cimport cython
from cython.parallel cimport parallel, prange

# Define a function pointer to a metric
ctypedef double (*metric)(double[: ,::1], np.intp_t, np.intp_t)

cdef extern from "math.h" nogil:
    double cos(double x)
    double sqrt(double x)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cdef double dihedral_distances(double[:, ::1] a, np.intp_t i1, np.intp_t i2):
    cdef double res
    cdef int m
    cdef int j

    res = 0.
    m = a.shape[1]

    for j in range(m):
        res += 1. - cos(a[i1, j] - a[i2, j])

    res /= 2.*m

    return sqrt(res)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cdef double dihedral_distances_p(double[:, ::1] a, np.intp_t i1, np.intp_t i2):
    cdef double res
    cdef int m
    cdef int j

    res = 0.
    m = a.shape[1]

    with nogil, parallel(num_threads=2):
        for j in prange(m, schedule='dynamic'):
            res += 1. - cos(a[i1, j] - a[i2, j])

    res /= 2.*m

    return sqrt(res)

@cython.boundscheck(False)
@cython.wraparound(False)
def pairwise(double[: ,::1] c not None, np.intp_t x, p = True):
    cdef metric dist_func
    if p:
        dist_func = &dihedral_distances_p
    else:
        dist_func = &dihedral_distances

    cdef np.intp_t i, n_structures
    n_samples = c.shape[0]

    cdef double[::1] res = np.empty(n_samples)

    for i in range(n_samples):
        res[i] = dist_func(c, x, i)

    return res

%timeit pairwise(c, x, False)
100 loops, best of 3: 17 ms per loop    

# Parallel version
%timeit pairwise(c, x, True)
10 loops, best of 3: 37.1 ms per loop

所以我跟着你的链接创建了dihedral distances函数的Cython版本。我们获得了一些速度提升，虽然不是很多，但仍然比numexpr版本慢（17毫秒对9.44毫秒）。所以我尝试使用prange并行化该函数，但结果更差（37.1毫秒对17毫秒对9.4毫秒）！

我错过了什么吗？

- NoExiT

3

几个小的改进是：1）将*0.5放在求和符号外面；2）在从1中减去之前对cos进行求和（这样做更准确，因为总和会接近于1）。对我来说，这些改进将运行时间从25毫秒缩短到了17毫秒。我知道你正在寻找更多的改进，但这是我的全部建议，希望能有所帮助。 - tom10

1

尝试使用Cython很容易，并且可以显著提高速度（当然，结果因人而异）：https://jakevdp.github.io/blog/2012/08/08/memoryview-benchmarks/ - ev-br

@tom10，我的程序虽然能够有效地获得1毫秒的时间，但是却出现了错误：“RuntimeWarning: invalid value encountered in sqrt”。 - NoExiT

2

对于2，如果你对N个项目求和，那么sqrt(sum(1 - cos(x)))就变成了sqrt(N - sum(cos(x)))。你记得N了吗？ - tom10

@tom10 是的，你说得对！我之前没有发现这个技巧。谢谢！ - NoExiT

2个回答

2

以下是使用Cython快速尝试的简单示例，仅涉及一对1D数组：

（在IPython笔记本中）

%%cython

cimport cython
cimport numpy as np

cdef extern from "math.h":
    double cos(double x) nogil
    double sqrt(double x) nogil

def cos_norm(a, b):
    return cos_norm_impl(a, b)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cdef double cos_norm_impl(double[::1] a, double[::1] b) nogil:
    cdef double res = 0., val
    cdef int m = a.shape[0]
    # XXX: shape of b not checked
    cdef int j

    for j in range(m):
        val = a[j] - b[j]
        res += 1. - cos(val)
    res /= 2.*m

    return sqrt(res)

与直接使用numpy实现相比，

def np_cos_norm(a, b):
    val = np.add.reduce(1. - np.cos(a-b))
    return np.sqrt(val / 2. / a.shape[0])

我明白了

np.random.seed(1234)

for n in [100, 1000, 10000]:
    x = np.random.random(n)
    y = np.random.random(n)
    %timeit cos_norm(x, y)
    %timeit np_cos_norm(x, y)
    print '\n'

100000 loops, best of 3: 3.04 µs per loop
100000 loops, best of 3: 12.4 µs per loop

100000 loops, best of 3: 18.8 µs per loop
10000 loops, best of 3: 30.8 µs per loop

1000 loops, best of 3: 196 µs per loop
1000 loops, best of 3: 223 µs per loop

因此，根据向量的维度，您可以获得4倍至无的加速。

对于计算成对距离，您可能可以做得更好，如这篇博客文章所示，但当然取决于您自己的情况。

- ev-br

谢谢@ev-br！我正在处理这个问题。np_cos_norm函数中只有一个小错误，应该是val = np.add.reduce(1. - np.cos(a-b), axis=1)。 - NoExiT

这是有意为之的：我只处理一维数组。 - ev-br

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

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

- serge-sans-paille · Accepted Answer

如果您愿意使用http://pythran.readthedocs.io/，您可以利用numpy实现，在这种情况下获得比cython更好的性能：

#pythran export np_cos_norm(float[], float[])
import numpy as np
def np_cos_norm(a, b):
    val = np.sum(1. - np.cos(a-b))
    return np.sqrt(val / 2. / a.shape[0])

并使用以下命令进行编译：

pythran fast.py

为了获得Cython版本的平均x2。

如果使用：

pythran fast.py -march=native -DUSE_BOOST_SIMD -fopenmp

您将获得一个矢量化的、并行的版本，运行速度略快：

100000 loops, best of 3: 2.54 µs per loop
1000000 loops, best of 3: 674 ns per loop

100000 loops, best of 3: 16.9 µs per loop
100000 loops, best of 3: 4.31 µs per loop

10000 loops, best of 3: 176 µs per loop
10000 loops, best of 3: 42.9 µs per loop

使用与 ev-br 相同的测试平台