Home编程语言python高性能的Python扩展（1）

高性能的Python扩展（1）

2017 年 1 月 1 日
本文由 伯乐在线 – Janzou 翻译。未经许可，禁止转载！
英文出处：J. David Lee。欢迎加入翻译组。

简介

通常来说，Python不是一种高性能的语言，在某种意义上，这种说法是真的。但是，随着以Numpy为中心的数学和科学软件包的生态圈的发展，达到合理的性能不会太困难。

当性能成为问题时，运行时间通常由几个函数决定。用C重写这些函数，通常能极大的提升性能。

在本系列的第一部分中，我们来看看如何使用NumPy的C API来编写C语言的Python扩展，以改善模型的性能。在以后的文章中，我们将在这里提出我们的解决方案，以进一步提升其性能。

文件

这篇文章中所涉及的文件可以在Github上获得。

模拟

作为这个练习的起点，我们将在像重力的力的作用下为N体来考虑二维N体的模拟。

以下是将用于存储我们世界的状态，以及一些临时变量的类。

Python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
    # lib/sim.py
 
    class World(object):
        “””World is a structure that holds the state of N bodies and
        additional variables.
 
        threads : (int) The number of threads to use for multithreaded
                  implementations.
 
        STATE OF THE WORLD:
 
        N : (int) The number of bodies in the simulation.
        m : (1D ndarray) The mass of each body.
        r : (2D ndarray) The position of each body.
        v : (2D ndarray) The velocity of each body.
        F : (2D ndarray) The force on each body.
 
        TEMPORARY VARIABLES:
 
        Ft : (3D ndarray) A 2D force array for each thread’s local storage.
        s  : (2D ndarray) The vectors from one body to all others.
        s3 : (1D ndarray) The norm of each s vector.
 
        NOTE: Ft is used by parallel algorithms for thread-local
              storage. s and s3 are only used by the Python
              implementation.
        “””
        def __init__(self, N, threads=1,
                     m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e–3):
            self.threads = threads
            self.N  = N
            self.m  = np.random.uniform(m_min, m_max, N)
            self.r  = np.random.uniform(–r_max, r_max, (N, 2))
            self.v  = np.random.uniform(–v_max, v_max, (N, 2))
            self.F  = np.zeros_like(self.r)
            self.Ft = np.zeros((threads, N, 2))
            self.s  = np.zeros_like(self.r)
            self.s3 = np.zeros_like(self.m)
            self.dt = dt

在开始模拟时，N体被随机分配质量m，位置r和速度v。对于每个时间步长，接下来的计算有：

  1. 合力F，每个体上的合力根据所有其他体的计算。
  2. 速度v，由于力的作用每个体的速度被改变。
  3. 位置R，由于速度每个体的位置被改变。

第一步是计算合力F，这将是我们的瓶颈。由于世界上存在的其他物体，单一物体上的力是所有作用力的总和。这导致复杂度为O（N^2）。速度v和位置r更新的复杂度都是O（N）。

如果你有兴趣，这篇维基百科的文章介绍了一些可以加快力的计算的近似方法。

纯Python

在纯Python中，使用NumPy数组是时间演变函数的一种实现方式，它为优化提供了一个起点，并涉及测试其他实现方式。

Python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
    # lib/sim.py
 
    def compute_F(w):
        “””Compute the force on each body in the world, w.”””
        for i in xrange(w.N):
            w.s[:] = w.r – w.r[i]
            w.s3[:] = (w.s[:,0]**2 + w.s[:,1]**2)**1.5
            w.s3[i] = 1.0 # This makes the self-force zero.
            w.F[i] = (w.m[i] * w.m[:,None] * w.s / w.s3[:,None]).sum(0)
 
    def evolve(w, steps):
        “””Evolve the world, w, through the given number of steps.”””
        for _ in xrange(steps):
            compute_F(w)
            w.v += w.F * w.dt / w.m[:,None]
            w.r += w.v * w.dt

合力计算的复杂度为O（N^2）的现象被NumPy的数组符号所掩盖。每个数组操作遍历数组元素。

可视化

这里是7个物体从随机初始状态开始演化的路径图：

State Machine Diagram

性能

为了实现这个基准，我们在项目目录下创建了一个脚本，包含如下内容：

Python

1
2
3
    import lib
    w = lib.World(101)
    lib.evolve(w, 4096)

我们使用cProfile模块来测试衡量这个脚本。

Python

1
    python –m cProfile –scum bench.py

前几行告诉我们，compute_F确实是我们的瓶颈，它占了超过99％的运行时间。

Python

1
2
3
4
5
6
7
8
9
10
11
    428710 function calls (428521 primitive calls) in 16.836 seconds
 
    Ordered by: cumulative time
 
    ncalls  tottime  percall  cumtime  percall filename:lineno(function)
         1    0.000    0.000   16.837   16.837 bench.py:2(<module>)
         1    0.062    0.062   16.756   16.756 sim.py:60(evolve)
      4096   15.551    0.004   16.693    0.004 sim.py:51(compute_F)
    413696    1.142    0.000    1.142    0.000 {method ‘sum’ ...
         3    0.002    0.001    0.115    0.038 __init__.py:1(<module>)
       ...

在Intel i5台式机上有101体，这种实现能够通过每秒257个时间步长演化世界。

简单的C扩展 1

在本节中，我们将看到一个C扩展模块实现演化的功能。当看完这一节时，这可能帮助我们获得一个C文件的副本。文件src/simple1.c，可以在GitHub上获得。

关于NumPy的C API的其他文档，请参阅NumPy的参考。Python的C API的详细文档在这里。

样板

文件中的第一件事情是先声明演化函数。这将直接用于下面的方法列表。

Python

1
    static PyObject *evolve(PyObject *self, PyObject *args);

接下来是方法列表。

Python

1
2
3
4
    static PyMethodDef methods[] = {
       { “evolve”, evolve, METH_VARARGS, “Doc string.”},
       { NULL, NULL, 0, NULL } /* Sentinel */
    };

这是为扩展模块的一个导出方法列表。这只有一个名为evolve方法。

样板的最后一部分是模块的初始化。

Python

1
2
3
4
    PyMODINIT_FUNC initsimple1(void) {
       (void) Py_InitModule(“simple1”, methods);
       import_array();
    }

另外，正如这里显示，initsimple1中的名称必须与Py_InitModule中的第一个参数匹配。对每个使用NumPy API的扩展而言，调用import_array是有必要的。

数组访问宏

数组访问的宏可以在数组中被用来正确地索引，无论数组被如何重塑或分片。这些宏也使用如下的代码使它们有更高的可读性。

Python

1
2
3
4
5
6
7
8
#define m(x0) (*(npy_float64*)((PyArray_DATA(py_m) +
(x0) * PyArray_STRIDES(py_m)[0])))
#define m_shape(i) (py_m->dimensions[(i)])
 
#define r(x0, x1) (*(npy_float64*)((PyArray_DATA(py_r) +
(x0) * PyArray_STRIDES(py_r)[0] +
(x1) * PyArray_STRIDES(py_r)[1])))
#define r_shape(i) (py_r->dimensions[(i)])

在这里，我们看到访问宏的一维和二维数组。具有更高维度的数组可以以类似的方式被访问。

在这些宏的帮助下，我们可以使用下面的代码循环r:

Python

1
2
3
4
5
for(i = 0; i < r_shape(0); ++i) {
    for(j = 0; j < r_shape(1); ++j) {
        r(i, j) = 0; // Zero all elements.
    }
}

命名标记

上面定义的宏，只在匹配NumPy的数组对象定义了正确的名称时才有效。在上面的代码中，数组被命名为py_m和py_r。为了在不同的方法中使用相同的宏，NumPy数组的名称需要保持一致。

计算力

特别是与上面五行的Python代码相比，计算力数组的方法显得颇为繁琐。

Python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
    static inline void compute_F(npy_int64 N,
                                 PyArrayObject *py_m,
                                 PyArrayObject *py_r,
                                 PyArrayObject *py_F) {
      npy_int64 i, j;
      npy_float64 sx, sy, Fx, Fy, s3, tmp;
    
      // Set all forces to zero.
      for(i = 0; i < N; ++i) {
        F(i, 0) = F(i, 1) = 0;
      }
    
      // Compute forces between pairs of bodies.
      for(i = 0; i < N; ++i) {
        for(j = i + 1; j < N; ++j) {
          sx = r(j, 0) – r(i, 0);
          sy = r(j, 1) – r(i, 1);
    
          s3 = sqrt(sx*sx + sy*sy);
          s3 *= s3 * s3;
    
          tmp = m(i) * m(j) / s3;
          Fx = tmp * sx;
          Fy = tmp * sy;
    
          F(i, 0) += Fx;
          F(i, 1) += Fy;
    
          F(j, 0) -= Fx;
          F(j, 1) -= Fy;
        }
      }
    }

请注意，我们使用牛顿第三定律（成对出现的力大小相等且方向相反）来降低内环范围。不幸的是，它的复杂度仍然为O（N^2）。

演化函数

该文件中的最后一个函数是导出的演化方法。

Python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    static PyObject *evolve(PyObject *self, PyObject *args) {
    
      // Declare variables.
      npy_int64 N, threads, steps, step, i;
      npy_float64 dt;
      PyArrayObject *py_m, *py_r, *py_v, *py_F;
    
      // Parse arguments.
      if (!PyArg_ParseTuple(args, “ldllO!O!O!O!”,
                            &threads,
                            &dt,
                            &steps,
                            &N,
                            &PyArray_Type, &py_m,
                            &PyArray_Type, &py_r,
                            &PyArray_Type, &py_v,
                            &PyArray_Type, &py_F)) {
        return NULL;
      }
    
      // Evolve the world.
      for(step = 0;  step< steps; ++step) {
        compute_F(N, py_m, py_r, py_F);
 
        for(i = 0; i < N; ++i) {
          v(i, 0) += F(i, 0) * dt / m(i);
          v(i, 1) += F(i, 1) * dt / m(i);
    
          r(i, 0) += v(i, 0) * dt;
          r(i, 1) += v(i, 1) * dt;
        }
      }
    
      Py_RETURN_NONE;
    }

在这里，我们看到了Python参数如何被解析。在该函数底部的时间步长循环中，我们看到的速度和位置向量的x和y分量的显式计算。

性能

C版本的演化方法比Python版本更快，这应该不足为奇。在上面提到的相同的i5台式机中，C实现的演化方法能够实现每秒17972个时间步长。相比Python实现，这方面有70倍的提升。

观察

注意，C代码一直保持尽可能的简单。输入参数和输出矩阵可以进行类型检查，并分配一个Python装饰器函数。删除分配，不仅能加快处理，而且消除了由Python对象不正确的引用计数造成的内存泄露（或更糟）。

下一部分

在本系列文章的下一部分，我们将通过发挥C-相邻NumPy矩阵的优势来提升这种实现的性能。之后，我们来看看使用英特尔的SIMD指令和OpenMP来进一步推进。

如果您有任何疑问，意见，建议或更正，请通过联系链接告诉我。

1 赞
收藏

1 评论





关于作者：Janzou


做一直倦懒的小猫咪，享受午后阳光（新浪微博：@Jan_zou）

个人主页 ·
我的文章

· 13

转载自演道,想查看更及时的互联网产品技术热点文章请点击http://go2live.cn

Related Posts

About The Author

bjmayor

程序员，码农，php,python,ios,android,go，产品经理，创业。