python擅长处理密集型cpu计算_Python里CPU密集型任务提速试验
众所周知，纯Python跑起来真的很慢。那得有多慢呢？本篇文章我们以mandelbrot set为例，来探讨一下Python在CPU密集型任务的运行速度，以及优化方法。
Mandelbrot set
网上查了查，好像有翻译为曼德勃罗集合，曼德布洛特集合等等。大概意思就是一堆满足如下条件的复数集合。
并且
理不理解这个概念其实没那么重要，反正就是要用Python运行下面这项任务。
def compute_mandelbrot_py(N_max, some_threshold, nx, ny):
# A grid of c-values
xs = np.linspace(-2, 1, nx)
ys = np.linspace(-1.5, 1.5, ny)
mandelbrot_set = np.zeros((len(xs), len(ys)))
for i in range(len(xs)):
for k in range(len(ys)):
x = xs[i]
y = ys[k]
c = x + 1j*y
z = c
is_inside = 1
for j in range(N_max):
z = z**2 + c
if abs(z) >= some_threshold:
is_inside = 0
break
mandelbrot_set[i, k] = is_inside
return mandelbrot_set
为了更好地观察和比较运行速度，我们把N_max与some_threshold统一分别取500和4。nx和ny分别代表了实数和虚数部分的样本，总共我们会计算nx * ny个复数来分别检验他们在不在这个集合里。
先取nx = 500, ny = 500来看看python会花多少时间。
电脑配置
Processor: 2.9 GHz Intel Core i7
Memory: 16 GB 2133 MHz LPDDR3
Graphics: Intel HD Graphics 630 1536 MB
开始运行
mandelbrot_set = compute_mandelbrot_py(500, 4., 500, 500)
花了9.8秒，嗯，真的是很慢，250000个数据点居然就要差不多10秒了，那1000 * 1000岂不是要炸?
尝试了下，果然是x4，花了41秒。那开始想办法提速吧。
1. 用Numpy把运算向量化
我们都知道Numpy向量运算代替循环可以优化计算，上面的头两个for loop因为里面的东西不取决于上一次循环的结果，所以自然的可以用矩阵表示并进行向量化运算。
c = x[:,np.newaxis] + 1j*y[np.newaxis,:]
z = c
for j in range(N_max):
z = z**2 + c
mandelbrot_set = (abs(z) < some_threshold)
这里的唯一缺点是不能像之前提前离开for j in range(N_max): 这个循环，但是结果显示还是比纯Python暴力循环方法快了不少，500 500花了0.47秒，1000 * 1000花了4.4秒。
2. 把密集运算部分放到Cython执行
Cython是结合了Python和C的语法的一种语言，可以简单认为就是给Python加上了静态类型后的语法，由于会直接编译成二进制程序，所以性能较Python而言会有很大提升。
我们先把mandelbrot运算部分的代码用cython编写然后方法一个pyx后缀的文件里。
cpdef mandelbrot(
double [:] xs,
double [:] ys,
double [:, :] mandelbrot_set,
int N_max,
int some_threshold):
cdef:
int i, k, j, is_inside
double x, y
complex c, z
for i in range(len(xs)):
for k in range(len(ys)):
x = xs[i]
y = ys[k]
c = x + 1j*y
z = c
is_inside = 1
for j in range(N_max):
z = z**2 + c
if abs(z) >= some_threshold:
is_inside = 0
break
mandelbrot_set[i, k] = is_inside
和纯Python比，这个代码把一些变量都加上了静态类型，我们把这个文件命名为mandelbrot_pyx.pyx。接下来要建立一个setup.py文件
from distutils.core import setup
from Cython.Build import cythonize
import numpy
setup(
ext_modules = cythonize("mandelbrot_pyx.pyx"),
include_dirs=[numpy.get_include()]
)
保存好这个文件后就可以开始编译cython代码了
python setup.py build_ext --inplace
这样在我们的python主程序里可以调用mandelbrot_pyx这个模块以及里面的mandelbrot函数。
mandelbrot_set = np.zeros((len(xs), len(ys)))
mandelbrot_pyx.mandelbrot(xs, ys, mandelbrot_set, N_max, some_threshold)
return mandelbrot_set
带入M = 1000, N = 1000， 发现只花了0.51秒，这比纯Python解法已经足足快了80倍了。于是我们再次增加M和N的值，让M = N = 2000，结果运算花了2.03秒，可以说是线性增加。
Cython里还有prange可以支持并行计算，在执行时得释放全局解释锁，可以有效实现多线程并行运算，把循环改成。
for i in prange(m, schedule="static", chunksize=1, num_threads=4, nogil=True):
如果只在本地跑的话随着M和N增大，运行速度还是会有点慢。那么怎么才能继续大幅度提高这个代码的效率呢？下面就来试试GPU并行吧。
3. GPU并行计算
由于我的电脑显卡并不是nvdia，所以这里不用CUDA，来试试opencl好了。
python部分，首先是启动opencl前的一些准备。
import opencl as cl
context = cl.Context(cl.get_platforms()[0].get_devices())
queue = cl.CommandQueue(context, context.devices[0])
program = cl.Program(context, open('mandelbrot.cl').read()).build(options='')
然后把input拷贝到gpu内存
gpu_real = cl.Buffer(context, cl.mem_flags.READ_ONLY, len(xs) * 4)
gpu_imag = cl.Buffer(context, cl.mem_flags.READ_ONLY, len(ys) * 4)
gpu_output = cl.Buffer(context, cl.mem_flags.WRITE_ONLY, len(xs) * len(ys) * 4)
cl.enqueue_copy(queue, gpu_real, xs, is_blocking=False)
cl.enqueue_copy(queue, gpu_imag, ys, is_blocking=False)
接下来是opencl部分，其实语法就是c语言的语法
__kernel void
mandelbrot(__global __read_only float *coords_real,
__global __read_only float *coords_imag,
__global __write_only int *mandelbrot_set,
int w, int h, int max_iter)
{
const int x = get_global_id(0);
const int y = get_global_id(1);
float c_real, c_imag;
float z_real, z_imag;
int iter;
if ((x < h) && (y < w)) {
float c_real = coords_real[x], c_imag = coords_imag[y];
float z_real = c_real, z_imag = c_imag;
bool is_inside = true;
for(unsigned i=0; i
float z_real2 = z_real * z_real, z_imag2 = z_imag * z_imag;
if (z_real2 + z_imag2 >= 4) {
is_inside = false;
break;
}
z_imag = 2 * z_real * z_imag + c_imag;
z_real = z_real2 - z_imag2 + c_real;
}
if (is_inside) {
mandelbrot_set[x * w + y] = 1;
}
}
}
把cl部分命名为mandelbrot.cl后，下面这句话就把这段代码编译了。
cl.Program(context, open('mandelbrot.cl').read()).build(options='')
把代码跑起来
program.mandelbrot(....)
再通过以下代码把结果从gpu拷贝回来。
cl.enqueue_copy(queue, mandelbrot_set, gpu_counts, is_blocking=True)
惊喜地发现2000 * 2000的集合代如运算只花了不到0.2秒！
总结
以上就是加速Python运行效率的一些小试验，并不一定适用于所有CPU密集型任务，但是希望有一定借鉴意义。下面附上一张运行速度的折线图。