[翻译]教程:当Numpy不够快时。。。

教程:当Numpy不够快时。。。原文链接

在python程序里使用fortran/blas提高6倍速度的教程

总结:一个demo用来展示怎么在python程序的矢量操作中使用fortran/blas库,来获得比Numpy更快的速度。

昨天,我从笔记本中发布了一篇怎么在ApacheSpark里使用GPU的文章。让我高兴的是,它进入了Hacker News的首页(而且是在服务Scala社区的情况下!!!)。在我的iPython notebook之一中使用Spark已经成为我的一份真实的热情。。。然而昨天我集中精力研究Scala/JVM/GPU的操作,今天我想提供给Python科学计算社区一些东西。这些发现来自于学习Radim Rehurek写的一个叫做Gensim的精彩的代码库。。。具体是word2vec的实现。

你可能会好奇,为什么我会在基于GPU加速的文章后面涉及基于CPU的加速。。。事实是有时更加轻量的优化更加合适。。。尤其是在处理更小批量的矢量或者GPU不可用时。

第一部分:iPython-Notebook Cython Magic

安装Numpy,ScipyGFortranCython,和scikit-learn包。我__强烈__建议坚持使用easy_install, brew(apt-get)和pip。以我的经验来看,macports在安装这些库时存在一些问题。还有,当然,你需要安装ipython notebook来使用这些例子,但从原理上来说,普通的cython也能使用。

使用Cython你也将得到“Cython”的魔法。下面的命令应该在你的notebook下有效。

%load_ext cythonmagic

%%cython
def add(int a, int b):
    return 3.1415926 * a + 2.718281828459045 * b

add(2, 2)

注意你使用%load_ext cythonmagic载入cython magic,然后在包含cython的代码块的顶部使用%%cython编译cython。然后你可以从python中调用你的cython函数(或者类。。等等)。这是个优雅的系统。:)

第二部分:Cython中的Scipy Fortran-Blas

下面你将看到我们超快的blas操作的核心代码。在开头的一些imports后,你会看到一个“cdef extern from”,引入了一个叫做voidptr.h的文件。这个文件可以让我们获得一个numpy数组的指针,而不需要拷贝任何数据。。。是代码的关键部分。文件的内容如下:

%%cython --annotate

import cython
import numpy as np
cimport numpy as np

from libc.math cimport exp
from libc.string cimport memset

from scipy.linalg.blas import fblas

REAL = np.float32
ctypedef np.float32_t REAL_t

cdef extern from "/Users/myname/Laboratory/voidptr.h":
    void* PyCObject_AsVoidPtr(object obj)


ctypedef void (*scopy_ptr) (const int *N, const float *X, const int *incX,
        float *Y, const int *incY) nogil
ctypedef void (*saxpy_ptr) (const int *N, const float *alpha, const float *X, const int *incX, float *Y, const int *incY) nogil
ctypedef float (*sdot_ptr) (const int *N, const float *X, const int *incX, const float *Y, const int *incY) nogil
ctypedef double (*dsdot_ptr) (const int *N, const float *X, const int *incX, const float *Y, const int *incY) nogil
ctypedef double (*snrm2_ptr) (const int *N, const float *X, const int *incX) nogil
ctypedef void (*sscal_ptr) (const int *N, const float *alpha, const float *X, const int *incX) nogil

cdef scopy_ptr scopy=<scopy_ptr>PyCObject_AsVoidPtr(fblas.scopy._cpointer)  # y = x
cdef saxpy_ptr saxpy=<saxpy_ptr>PyCObject_AsVoidPtr(fblas.saxpy._cpointer)  # y += alpha * x
cdef sdot_ptr sdot=<sdot_ptr>PyCObject_AsVoidPtr(fblas.sdot._cpointer)  # float = dot(x, y)
cdef dsdot_ptr dsdot=<dsdot_ptr>PyCObject_AsVoidPtr(fblas.sdot._cpointer)  # double = dot(x, y)
cdef snrm2_ptr snrm2=<snrm2_ptr>PyCObject_AsVoidPtr(fblas.snrm2._cpointer)  # sqrt(x^2)
cdef sscal_ptr sscal=<sscal_ptr>PyCObject_AsVoidPtr(fblas.sscal._cpointer) # x = alpha * x

cdef int ONE = 1
cdef REAL_t ONEF = <REAL_t>1.0

def putDotty(syn0, syn1, size):
    cdef int lSize = size
    f = <REAL_t>sdot(&lSize, <REAL_t *>(np.PyArray_DATA(syn0)), &ONE, <REAL_t *>(np.PyArray_DATA(syn1)), &ONE)
    return f

voidptr.h代码,地址见github

#include <Python.h>

#if PY_VERSION_HEX >= 0x03020000

/*
** compatibility with python >= 3.2, which doesn't have CObject anymore
*/
static void * PyCObject_AsVoidPtr(PyObject *obj)
{
    void *ret = PyCapsule_GetPointer(obj, NULL);
    if (ret == NULL) {
        PyErr_Clear();
    }
    return ret;
}

#endif

下面,你可以看到六个函数类型和它们的实现。在Scipy Blas/Fortran文档里有一整套这些时髦命名的fortran函数。我也写了一个简单的点乘函数pubDotty,利用了dsdot函数(double dot product。。和float相对)。(译者注:我使用dsdot结果为0,改为sdot结果正常。)

np.random.seed(0)
x = np.zeros(32, dtype='f')
y = np.zeros(32, dtype='f')

x += 1
y += 3

%timeit np.dot(x, y)
1000000 loops, best of 3: 1.33 us per loop

%timeit pubDotty(x, y, len(x))
1000000 loops, best of 3: 226 ns per loop

在这个例子中,我创建了两个32长度的numpy矢量。(一个全是1,另一个全是3)。然后我做了基准测试,cython/fortran版本快了5.8倍。应该注意这依然是传入了一个python对象。。。当所有高级操作都在cython中时,这个效率可以进一步提升。


我也用IPython notebook做个下这个实验(结果见下图),numpy用时508ns,而这个函数用时125ns,确实有4倍多的提升~~

FireShot Capture.png