cython通过编译为C程序提高性能有很多例子,通过OpenMP并行的性能没那么多。
今天尝试了一下似乎gcc对parallelism reduction优化的很厉害,加上OpenMP并行可以提高20倍性能,这不是简单的2 core带来的性能提高。
滚动求和 rolling sum的例子
最简单的实现pandas.rolling,通过操作numpy array,速度也还算能接受。
# test_para.py
import numpy as np
#import pyximport; pyximport.install(reload_support=True, setup_args={"include_dirs":np.get_include()})
import timeit
import pandas as pd
import para.cpara as cpara
X = -1 + 2*np.random.rand(100000)
ss = pd.Series(X)
ss.rolling(100).apply(np.sum,raw=True)
print('==============')
print('multi thread')
start_time = timeit.default_timer()
sum_cython=pd.Series(cpara.window_sum(X, 100))
print(timeit.default_timer() - start_time)
print('single thread')
start_time = timeit.default_timer()
sum_pandas=ss.rolling(100).apply(np.sum,raw=True)
print(timeit.default_timer() - start_time)
print(np.max(np.abs(sum_cython - sum_pandas)))
cython源文件
# cpara.pyx
cimport cython
import numpy as np
from cython.parallel import prange,parallel
cimport numpy as cnp
from libc.stdlib cimport malloc
@cython.boundscheck(False)
def window_sum(cnp.ndarray[double, ndim=1] arr, int window):
cdef h = np.zeros_like(arr)
cdef int imax = arr.shape[0]
cdef double *buffer = <double *>malloc(imax * sizeof(double))
cdef double result = 0.0
cdef int i, j
with nogil, parallel():
for i in prange(imax, schedule='dynamic'):
buffer[i] = 0.0
if i >= window-1:
for j in range(window):
buffer[i] += arr[i-j]
for i in range(imax):
if i < window -1:
h[i] = np.nan
else:
h[i] = buffer[i]
return h
setup.py中要加入openmp的编译链接参数
EXT = Extension("*",
["para/*.pyx"],
define_macros=[('CYTHON_TRACE', CYTHON_DEBUG),
('CYTHON_TRACE_NOGIL', CYTHON_DEBUG),
('CYTHON_BINDING', CYTHON_DEBUG),
("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION"),
('CYTHON_FAST_PYCCALL', '1')],
extra_compile_args = ["-fopenmp" ],
extra_link_args=['-fopenmp'],
include_dirs=[".", np.get_include()])
性能比较
%timeit pd.Series(cpara.window_sum(X, 100))
23.4 ms ± 325 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit ss.rolling(100).apply(np.sum,raw=True)
536 ms ± 3.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
536/23.4=22.9
测试环境:i3-7100U 2core 2T CPU, ubuntu 18.04 LTS