Tối ưu hóa và tăng tốc của một hàm toán học trong python

Question

Mục đích của chức năng toán học này là để tính toán khoảng cách giữa hai (hoặc nhiều) cấu trúc protein sử dụng những góc nhị diện:

Nó rất hữu ích trong việc cấu trúc sinh học, ví dụ. Và tôi đã mã hóa hàm này trong python bằng cách sử dụng numpy, nhưng mục đích là để có một triển khai nhanh hơn. Là tham chiếu thời gian tính toán, tôi sử dụng hàm khoảng cách euclide có sẵn trong gói tìm hiểu scikit.

Dưới đây là đoạn code tôi có cho thời điểm này:

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 ms nó rất nhanh, nhưng nó rất chậm nếu bạn cần phải chạy nó một triệu lần. Bây giờ câu hỏi là, làm thế nào để làm điều đó? Bước tiếp theo là gì? Cython? PyOpenCL? Tôi có một số kinh nghiệm với PyOpenCL, tuy nhiên tôi không bao giờ viết một cái gì đó phức tạp như thế này. Tôi không biết nếu nó có thể tính toán khoảng cách dihedral trong một bước trên GPU như tôi làm với numpy và làm thế nào để tiến hành.

Cảm ơn bạn đã giúp tôi!

EDIT: Cảm ơn các bạn! Tôi hiện đang làm việc trên các giải pháp đầy đủ và một khi nó đã hoàn thành, tôi sẽ đặt mã ở đây.

Cython VERSION:

%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 

Vì vậy, tôi làm theo bạn link để tạo ra các phiên bản cython của quãng đường Hàm nhị diện. Chúng tôi đạt được một số tốc độ, không quá nhiều, nhưng nó vẫn còn chậm hơn so với phiên bản numexpr (17ms so với 9.44ms). Vì vậy, tôi đã cố gắng để song song chức năng bằng cách sử dụng prange và nó là tồi tệ hơn (37.1ms vs 17ms vs 9.4ms)!

Tôi có bỏ lỡ điều gì đó không?

Answer 1

Nếu bạn sẵn sàng để sử dụng http://pythran.readthedocs.io/, bạn có thể tận dụng tình hình thực hiện NumPy và có được hiệu suất tốt hơn so với cython cho trường hợp đó:

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

Và biên dịch nó với:

pythran fast.py 

Để nhận được x2 trung bình so với phiên bản cython.

Nếu sử dụng:

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

Bạn sẽ nhận được một vectorized, phiên bản song song chạy nhanh hơn một chút:

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 

(bằng cách sử dụng nền tảng thử nghiệm tương tự như ev-br)

Answer 2

Dưới đây là một thử nhanh chóng-và-bẩn với cython, chỉ một cặp mảng 1D:

(trong một máy tính xách tay 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) 

So sánh với thực hiện NumPy đơn giản,

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

Tôi nhận được

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 

Vì vậy, tùy thuộc vào kích thước của vectơ của bạn, bạn có thể nhận được từ hệ số 4 đến nil của một tăng tốc.

Để tính toán khoảng cách theo cặp, bạn có thể làm tốt hơn nhiều, như được hiển thị trong this blog post, nhưng tất nhiên là YMMV.