python高性能优化

经过一个月的努力，终于成功优化了我的python程序，记录下一路上踩过的坑

1.原始程序

我的原始程序是模拟一个随时间演化的多粒子系统，每个粒子跟它的邻居（每个粒子有各自的的邻居）有相互作用，考虑到我的论文还没发表，这里先不贴出代码。大致上，程序有一个大循环，用来做时间步的迭代；大循环里有个小循环，用来对N=91个粒子进行迭代。一轮要跑100万个时间步，一组实验要跑10轮，然后改变参数还要进行上百组实验。测试原始代码的性能大概是跑2000步用时30秒，这么算下来，用一个10核的服务器（可以同时跑10个进程），一次实验也得三个星期才能完成。这么慢的速度我的论文猴年马月才能写完啊！无奈之下开始踏上python高性能计算的不归路。

补个题外话，当初要是听我导师的话用C++写可能也没有这么多优化的问题了，python虽然开发比C容易的多，但性能真的不及C的百分之一。

2.多线程与多进程

一个很自然的想法是把代码改成多线程的，在小循环里开多个线程同时跑N个粒子，然后python的GIL线程锁直接否定了这个办法，对于计算密集型的程序，python多线程根本就是假的多线程。

既然线程不行，那就考虑多进程吧。python的多进程还算好写，导入multiprocessing，开一个进程池就ok了。然后代码一跑，比原始代码还慢？？分析了下应该是因为创建进程的开销太大了，100万步每一步都要创建进程池可想而知多费时。然后尝试了下把代码结构改变，让N个粒子分成几份各自独立的跑，这样只需要初始创建进程池而不用每一步都创建了，但是又有个问题：怎么同步数据？因为当前粒子的位置和速度是根据上一步的信息计算得到的，所以必须等所有粒子都更新完才能进行下一步的更新，这就需要每个进程跑完一个时间步后要等待其他进程也跑完，然后同步各自的数据开始下一步。于是花了好几天研究进程的同步与数据共享，然后……我从入门到放弃了，bug太多了(T_T)

另外，多进程受限于cpu的核数，最多只能加速cpu核数那么多倍。也就是说，我开10个进程跑一个程序，实际上跟我跑10个程序每个程序单进程是一样的。所以把代码改成多进程，实际上跟我把一组实验分成好几组在服务器上跑是异曲同工的。

3.纯python优化

要改进程序的性能，还是先从程序本身入手。纯粹的python优化包括选择合适的算法和数据结构。我改进了部分代码，把没用的计算过程都删了，数据全部改成用numpy数组计算，尽量避免了在循环中导入方法。一个很有效的优化方式是，把显示的for循环改成列表推导，[x for _ in range(N) ]这样子，然后求和直接用sum函数。优化后，性能提高了大概30%，原来2000步30秒，现在只用22秒。

4.Cython混编

代码优化提高30%的性能还远远不够啊，在网上查到一个python的C拓展，叫Cython。

俗话说，”你不能打败一个用C写的循环“。要是能让python拥有C一样的速度多好啊！

于是，把循环调出，定义一大堆C的数据结构，改用Cython写循环，保存成.pyx文件用gcc编译，然后bug，bug，bug，……ok，还好调通了，在服务器上跑测试程序，bug，bug，bug，……各种编译错误，重装gcc又是各种问题，于是，放弃了╮(╯▽╰)╭

5.numba加速

踩了这么多坑终于让我用上了numba，numba是一个加速库，可以直接将小型函数编译成机器码，运行速度极快。

numba的使用很简单，只需要在函数前面加一个jit装饰器就行了：

from numba import nb
@nb.jit
def update(particles):
    ...

要发挥numba的最大性能，需要给装饰器传递签名，用来指定输入和输出的数据类型，像这样：

@nb.jit(["void(float64[:,:], float64[:], int8[:,:], float64, float64)"], nopython=True)
def update(particles, thetas, nei_matrix, l, Da):
    ...

这里指定了函数返回值是void，particles是一个浮点二维数组用float64[:,:]表示，thetas是浮点一维数组float64[:]，其他类似。"nopython=True"是编译选项。

使用numba关键的一点是要规范好函数里的数据类型和操作，否则的话起不到加速效果甚至没法编译，比如函数里如果有字符串或者嵌套列表的操作，很可能引起编译错误。经过一系列数学推导，我把算法里的循环做了优化，然后所有的数据操作都改成基于“原生”的数组操作，没有嵌套也没有运用numpy广播，尽量少的调用方法，循环也都直接用显示的for循环。ok，这样编译顺利通过了。然后测试一下，2000步只用了0.28秒！速度提升了将近100倍！O(∩_∩)O

最后，numba实际上还支持GPU计算，要知道处理大规模数组操作，GPU是一大杀器。有时间用GPU优化试试。