提高Python程序的运行速度
尝试了一下用Python实现的K-Means Clustering算法,抽样了10000篇百科词条,分为1000个类,分词后词语总数为130000左右。如果把1000个类定义为1000个向量,每个向量的元素个数为130000,K-Means Clustering算法的第一步是初始化这1000个向量的值,如果每个向量元素的值用float型存储,则需要的内存为:
1000 * 130000 * sizeof(float)
约 520M 左右。
最初用Python的list存储,动态扩展列表大小,结果内存用到近3G还没初始化完成,只好赶紧kill掉了。
改用 NumPy 存成固定数组的类别,发现内存使用量和计算结果基本一致,而且NumPy支持数组的数学计算,确实方便了不少,但即便如此,性能仍不够理想,K-Means算法第一次迭代花了一个小时还没完成。
怀疑K-Means算法的pearson距离计算时间较长(没有用profile工具论证过),用 Cython重新改写了该算法,并尽可能缓存计算中间结果。第一轮迭代耗时1个多小时,算是有了进步。
由于K-Means算法费CPU较多,改写了计算逻辑,用multiprocessing模块并行计算,在4核的CPU上速度提升了4倍,第一轮迭代花了约30分钟。但multiprocessing需要fork多个进程,每个进程的内存使用量均在600M上下想,4个进程占用了2.4G内存,代价有些大。而且计算结果在不同的进程间传递,性能开销也是存在的。 multiprocessing + 共享内存可以解决这个问题,但Python的数据结构不好表示。
继而想到写C的动态链接库,用Python的ctypes模块调用该动态链接库完成计算过程,C的动态链接库则创建系统线程,这样能有效躲过Python的GIL。问题是C代码有时候确实需要访问Python的数据对象,这只能通过Python的C扩展模块实现了,但C的扩展模块能访问ctypes的原始C指针吗?如果只能通过Python的C API访问,则GIL是绕不开的,我们的目标是尽可能少地锁住Python虚拟机。
解决办法是修改Python的ctypes源码,让它导出函数 addressof 的C API,这样在其他的C扩展模块里就能拿到ctypes的原始数据块指针。addressof的返回结果是一个long的PyObject包装对象,通过 PyLong_AsVoidPtr调用即可获取其值。
因此最后的解决办法是下载Python的源码,修改模块_ctypes的源码,通过Capsule导出C API。继而编写Python的C扩展模块,创建线程池,直接操作ctypes定义的数据内存。由于Python数据结构是非线程安全的,访问它们仍需获得GIL,但用到的可能性很小。程序的主体逻辑仍由Python代码构造,包括必要的ctypes数据结构。
代码改写后一次K-Means迭代不到一分钟就完成了,4颗CPU全跑满,内存仅占用600M左右,跟预期完全一致。
总结
Python大多数时间能如我们预期那样工作
涉及到数值计算和海量循环,Python表现极其糟糕
Cython + NumPy 可解决部分计算问题和内存问题,但GIL无法避免
multiprocessing能解决SMP/GIL问题,但内存问题解决不了,也许共享内存+ctypes是个办法,没尝试过
ctypes + C的动态链接库创建系统线程能解决GIL和内存共享问题,但无法在C中操作Python对象
ctypes + Python C扩展意义不大,因为C扩展无法直接操作ctypes的数据指针
改写后的ctypes + Python C扩展解决了性能和内存消耗问题
ctypes数据类型能被msgpack之类的工具快速序列化到磁盘,但恢复的时候只能恢复成list类型,这可能是美中不足的地方。