我刚刚编写了一个简单的程序来测试cython
的prange
的性能,下面是代码:
from cython.parallel import prange
import numpy as np
def func(int r, int c):
cdef:
double[:,:] a = np.arange(r*c, dtype=np.double).reshape(r,c)
double total = 0
int i, j
for i in prange(r, nogil=True, schedule='static', chunksize=1):
for j in range(c):
total += a[i,j]
return total
在使用
OMP_NUM_THREADS=3
的Mac Book pro上,上面的代码对于(r,c) = (10000, 100000)
大约需要18秒,而对于单线程,大约需要21秒。为什么性能提升这么少?我是否正确使用此
prange
? 最佳答案
您是否定时分配a
需要多长时间? 10000 x 100000 float64阵列占用8GB内存。
a = np.ones((10000, 100000), np.double)
配备16GB RAM的笔记本电脑需要6秒钟以上的时间。如果您没有8GB的可用空间,则需要进行交换,这将花费更长的时间。由于
func
几乎将所有时间都花在分配a
上,因此并行化外部for
循环只能在总运行时间上获得少量改进。为了说明这一点,我修改了您的函数以接受
a
作为输入。在tmp.pyx
中:#cython: boundscheck=False, wraparound=False, initializedcheck=False
from cython.parallel cimport prange
def serial(double[:, :] a):
cdef:
double total = 0
int i, j
for i in range(a.shape[0]):
for j in range(a.shape[1]):
total += a[i, j]
return total
def parallel(double[:, :] a):
cdef:
double total = 0
int i, j
for i in prange(a.shape[0], nogil=True, schedule='static', chunksize=1):
for j in range(a.shape[1]):
total += a[i, j]
return total
例如:
In [1]: import tmp
In [2]: r, c = 10000, 100000
In [3]: a = np.random.randn(r, c) # this takes ~6.75 sec
In [4]: %timeit tmp.serial(a)
1 loops, best of 3: 1.25 s per loop
In [5]: %timeit tmp.parallel(a)
1 loops, best of 3: 450 ms per loop
并行化该功能可以使我的4核笔记本电脑的速度提高2.8倍*,但这只是分配
a
所需时间的一小部分。这里的课程是始终分析您的代码,以了解在投入优化之前大部分时间都花在了哪里。
*通过将更大的
a
块传递给每个工作进程,您可以做得更好一些,例如通过增加chunksize
或使用schedule='guided'