使用并行Cython加速线性变换

时间:2017-04-03 22:44:49

标签: numpy parallel-processing cython numexpr

我需要加快线性变换的计算,该变换大致具有以下形式:

import numpy as np

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)
X,Y=np.meshgrid(x,y,sparse=True)

output=np.dot(np.cos(X*Y),input)

也就是说,我在常规网格上评估余弦并将输入乘以得到的矩阵。实际上,核函数(这里是余弦)更复杂,特别是它不是周期性的。因此,不可能简化FFT类型!

我的多核机器上面的转换大约需要5秒钟。现在,我肯定需要加快速度。首先尝试使用numexpr:

import numpy as np
import numexpr as ne

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)
X,Y=np.meshgrid(x,y,sparse=True)

output=np.dot(ne.evaluate('cos(X*Y)'),input)

这利用了并行计算,并将执行时间缩短到大约0.9秒。这很好,但不足以达到我的目的。所以,我的下一个尝试是采用并行Cython:

import numpy as np
from cython.parallel import prange

cimport numpy as np
cimport cython
from libc.math cimport cos

DTYPE = np.float64    
ctypedef np.float64_t DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def transform(double[:] x, double[:] y, double[:] input):

    cdef unsigned int N = x.shape[0]
    cdef double[:] output = np.zeros(N)
    cdef unsigned int row, col

    for row in prange(N, nogil= True):
        for col in range(N):
            output[row] += cos(x[row]*y[col])*input[col]

    return output

我通过执行

来编译它
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext

ext_modules=[
    Extension("cythontransform",
          ["cythontransform.pyx"],
          libraries=["m"],
          extra_compile_args = ["-O3", "-ffast-math", "-march=native", "-fopenmp" ],
          extra_link_args=['-fopenmp']
          ) 
]

setup(
  name = "cythontransform",
  cmdclass = {"build_ext": build_ext},
  ext_modules = ext_modules
)
从命令行

。通过

调用转换
import numpy as np
from cythontransform import transform

N=10000
input=np.random.random(N)

x=np.linspace(0,100,N)
y=np.linspace(0,30,N)

output=transform(x,y,input)

产生相当微弱的改善,大约0.7秒。

是否有人意识到进一步改进Cython代码的可能性?

或者,是否有其他一些框架(PyOpenCL,Pythran,Numba,...)更适合这个问题?

1 个答案:

答案 0 :(得分:1)

在我的笔记本电脑上,以下pythran版本:

#pythran export transform(float64[], float64[], float64[])
import numpy as np

def transform(x, y, input):
    N = x.shape[0]
    output = np.zeros(N)

    #omp parallel for
    for row in range(N):
        for col in range(N):
            output[row] += np.cos(x[row]*y[col])*input[col]
    return output

编译

pythran python -Ofast dd.py -fopenmp

运行速度大约是您提出的cython版本的两倍。我没有调查为什么会发生这种情况......