将Python加速器(Cython、Numba、f2py)与Numpy einsum进行比较

14
我正在比较Python加速器(Numba,Cython,f2py)与简单的For循环和Numpy的einsum,以解决特定问题(见下文)。到目前为止,对于这个问题,Numpy是最快的(因素快6倍),但我想知道是否有其他优化可以尝试,或者我做错了什么。这段简单的代码基于一个更大的代码,其中有许多这些einsum调用,但没有明确的for循环。我正在检查这些加速器中是否有任何一种可以做得更好。
使用Python 2.7.9在Mac OS X Yosemite上进行计时,使用Homebrew安装的gcc-5.3.0(--with-fortran --without-multilib)。还进行了%timeit调用;这些单个调用的时间相当准确。
In [1]: %run -i test_numba.py
test_numpy: 0.0805640220642
Matches Numpy output: True

test_dumb: 1.43043899536
Matches Numpy output: True

test_numba: 0.464295864105
Matches Numpy output: True

test_cython: 0.627640008926
Matches Numpy output: True

test_f2py: 5.01890516281
Matches Numpy output: True

test_f2py_order: 2.31424307823
Matches Numpy output: True

test_f2py_reorder: 0.507861852646
Matches Numpy output: True

主要代码:
import numpy as np
import numba
import time
import test_f2py as tf2py
import pyximport
pyximport.install(setup_args={'include_dirs':np.get_include()})
import test_cython as tcyth

def test_dumb(f,b):
    fnew = np.empty((f.shape[1],f.shape[2]))
    for i in range(f.shape[0]):
        for l in range(f.shape[3]):
            fnew += f[i,:,:,l] * b[i,l]
    return fnew


def test_dumber(f,b):
    fnew = np.empty((f.shape[1],f.shape[2]))
    for i in range(f.shape[0]):
        for j in range(f.shape[1]):
            for k in range(f.shape[2]):
                for l in range(f.shape[3]):
                    fnew[j,k] += f[i,j,k,l] * b[i,l]
    return fnew

@numba.jit(nopython=True)
def test_numba(f,b):
    fnew = np.zeros((f.shape[1],f.shape[2])) #NOTE: can't be empty, gives errors
    for i in range(f.shape[0]):
        for j in range(f.shape[1]):
            for k in range(f.shape[2]):
                for l in range(f.shape[3]):
                    fnew[j,k] += f[i,j,k,l] * b[i,l]
    return fnew

def test_numpy(f,b):
    return np.einsum('i...k,ik->...',f,b)

def test_f2py(f,b):
    return tf2py.test_f2py(f,b)

def test_f2py_order(f,b):
    return tf2py.test_f2py(f,b)

def test_f2py_reorder(f,b):
    return tf2py.test_f2py_reorder(f,b)

def test_cython(f,b):
    return tcyth.test_cython(f,b)

if __name__ == '__main__':

    #goal is to create: fnew = sum f*b over dim 0 and 3.
    f = np.random.rand(32,33,2000,64)
    b = np.random.rand(32,64)

    f1 = np.asfortranarray(f)
    b1 = np.asfortranarray(b)

    f2 = np.asfortranarray(np.transpose(f,[1,2,0,3]))

    funcs = [test_dumb,test_numba, test_cython, \
            test_f2py,test_f2py_order,test_f2py_reorder]

    tstart = time.time()    
    fnew_numpy= test_numpy(f,b)
    tstop = time.time()
    print test_numpy.__name__+': '+str(tstop-tstart)
    print 'Matches Numpy output: '+str(np.allclose(fnew_numpy,fnew_numpy))
    print ''

    for func in funcs:
        tstart = time.time()
        if func.__name__ == 'test_f2py_order':
            fnew = func(f1,b1)
        elif func.__name__ == 'test_f2py_reorder':
            fnew = func(f2,b1)
        else:
            fnew = func(f,b)
        tstop = time.time()
        print func.__name__+': '+str(tstop-tstart)
        print 'Matches Numpy output: '+str(np.allclose(fnew,fnew_numpy))
        print ''

这是一个使用f2py编译的文件(test_f2py.F90),命令为f2py -c -m test_f2py test_f2py.F90:

!file: test_f2py
subroutine test_f2py(f,b,fnew,n1,n2,n3,n4)

integer  :: n1,n2,n3,n4
real(8), dimension(n1,n2,n3,n4) :: f
real(8), dimension(n1,n4) :: b
real(8), dimension(n2,n3) :: fnew
!f2py intent(in) f
!f2py intent(in) b
!f2py intent(out) fnew
!f2py intent(in) n1
!f2py intent(in) n2
!f2py intent(in) n3
!f2py intent(in) n4

integer :: i1,i2,i3,i4

do i1=1,n1
    do i2=1,n2
        do i3=1,n3
            do i4=1,n4
                fnew(i2,i3) = fnew(i2,i3) + f(i1,i2,i3,i4)*b(i1,i4)
            enddo
        enddo
    enddo
enddo

end subroutine test_f2py

subroutine test_f2py_reorder(f,b,fnew,n1,n2,n3,n4)

integer  :: n1,n2,n3,n4
real(8), dimension(n1,n2,n3,n4) :: f
real(8), dimension(n3,n4) :: b
real(8), dimension(n1,n2) :: fnew
!f2py intent(in) f
!f2py intent(in) b
!f2py intent(out) fnew
!f2py intent(in) n1
!f2py intent(in) n2
!f2py intent(in) n3
!f2py intent(in) n4

integer :: i1,i2,i3,i4

do i3=1,n3
    do i4=1,n4
        do i1=1,n1
            do i2=1,n2
                fnew(i1,i2) = fnew(i1,i2) + f(i1,i2,i3,i4)*b(i3,i4)
            enddo
        enddo
    enddo
enddo

end subroutine test_f2py_reorder

而Cython .pyx文件(在主例程中使用pyximport编译):

#/usr/bin python
import numpy as np
cimport numpy as np

def test_cython(np.ndarray[np.float64_t,ndim=4] f, np.ndarray[np.float64_t,ndim=2] b):
    # cdef np.ndarray[np.float64_t,ndim=4] f
    # cdef np.ndarray[np.float64_t,ndim=2] b
    cdef np.ndarray[np.float64_t,ndim=2] fnew = np.empty((f.shape[1],f.shape[2]),dtype=np.float64)
    cdef int i,j,k,l
    cdef int Ni = f.shape[0]
    cdef int Nj = f.shape[1]
    cdef int Nk = f.shape[2]
    cdef int Nl = f.shape[3]

    for i in range(Ni):
        for j in range(Nj):
            for k in range(Nk):
                for l in range(Nl):
                    fnew[j,k] += f[i,j,k,l] * b[i,l]
    return fnew

由于您已经有可工作的代码,您的问题可能更适合于CodeReview.SE - ali_m
2
在我的笔记本电脑上(OSX 10.9.5),运行Numba 0.23.1的test_numpy()使用%timeit每个循环需要75.5毫秒,而test_numba()每个循环需要123毫秒,因此差异似乎不像您的测试那么极端。当对numba代码进行基准测试时,您要特别小心,确保至少调用一次以使代码在基准测试之外被jit编译,否则您将在数字中包含该成本,而每个后续调用将更快。 - JoshAdel
2个回答

11

通常这些加速器用于加速带有Python循环或大量中间结果的代码,而einsum已经非常优化(请参见源代码)。您不应该期望它们轻易地击败einsum,但在性能上可能会接近。

对于Numba来说,将编译时间从基准测试中排除是很重要的。这可以通过简单地运行两次已经被jit编译过的函数(使用相同类型的输入)来实现。例如,在IPython中我得到:

f = np.random.rand(32,33,500,64)
b = np.random.rand(32,64)

%time _ = test_numba(f,b)  # First invocation
# Wall time: 466 ms
%time _ = test_numba(f,b)
# Wall time: 73 ms
%timeit test_numba(f, b)
# 10 loops, best of 3: 72.7 ms per loop
%timeit test_numpy(f, b)
# 10 loops, best of 3: 62.8 ms per loop

对于您的Cython代码,可以进行一些改进:

  1. 禁用数组边界检查和环绕检查,请参见编译器指令
  2. 指定数组是连续的。
  3. 使用类型化内存视图

类似这样:

cimport cython
import numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
def test_cython(double[:,:,:,::1] f, double[:,::1] b):
    cdef int i, j, k, l, Ni, Nj, Nk, Nl
    Ni = f.shape[0]
    Nj = f.shape[1]
    Nk = f.shape[2]
    Nl = f.shape[3]

    fnew = np.empty((Nj, Nk))
    cdef double[:,::1] fnew_v = fnew

    for i in range(Ni):
        for j in range(Nj):
            for k in range(Nk):
                for l in range(Nl):
                    fnew_v[j,k] += f[i,j,k,l] * b[i,l]
    return fnew

在最新的Ubuntu 15.10(x86)上,这使我获得了与einsum相同的速度。然而,在Windows(x86)上,使用Anaconda发行版的同一台PC上,这个Cython代码的速度约为einsum的一半。我认为这可能与gcc版本(5.2.1 vs 4.7.0)以及插入SSE指令的能力有关(einsum使用SSE2内部函数编码)。也许提供不同的编译器选项会有所帮助,但我不确定。
我几乎不懂Fortran,所以无法发表评论。
既然你的目标是打败einsum,那么下一个明显的步骤就是增加并行性。使用cython.parallel可以很容易地生成一些线程。如果这还没有饱和您系统的内存带宽,那么您可以尝试显式包含最新的CPU指令,如AVX2和Fused Multiply-Add。

你可以尝试重新排列和重塑f,并使用np.dot进行操作。如果你的Numpy带有良好的BLAS库,这应该能够实现你所想到的几乎所有优化,但代价是失去了通用性,可能需要非常昂贵的f数组拷贝。


1
一旦解析字符串参数完成,einsum使用编译版本的nditer在所有轴上执行积和计算。源代码很容易在numpy github上找到。
不久前,我作为撰写补丁的一部分,研究了一个类似于einsum的工具。作为其中的一部分,我编写了一个cython脚本来完成积和运算。您可以在以下链接中查看此代码:

https://github.com/hpaulj/numpy-einsum

我没有试图让我的代码以einsum的速度运行。我只是想弄清楚它是如何工作的。


网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接