里的open函数_Python里CPU密集型任务提速试验-CSDN博客

众所周知，纯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<max_iter; ++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密集型任务，但是希望有一定借鉴意义。下面附上一张运行速度的折线图。