Python/Cython/Numpy 对 np.nonzero 的优化

Question

我有一段代码正在尝试优化。大部分代码执行时间由 cdef np.ndarray index = np.argwhere(array==1) 占用其中 array 是一个 numpy 是一个 512x512,512 numpy array of zeros and ones。关于加快速度的任何想法？使用 Python 2.7、Numpy 1.8.1

球度函数

def sphericity(self,array):

    #Pass an mask array (1's are marked, 0's ignored)
    cdef np.ndarray index = np.argwhere(array==1)
    cdef int xSize,ySize,zSize
    xSize,ySize,zSize=array.shape

    cdef int sa,vol,voxelIndex,x,y,z,neighbors,xDiff,yDiff,zDiff,x1,y1,z1
    cdef float onethird,twothirds,sp
    sa=vol=0 #keep running tally of volume and surface area
    #cdef int nonZeroCount = (array != 0).sum() #Replaces np.count_nonzero(array) for speed
    for voxelIndex in range(np.count_nonzero(array)):
    #for voxelIndex in range(nonZeroCount):
        x=index[voxelIndex,0]
        y=index[voxelIndex,1]
        z=index[voxelIndex,2]
        #print x,y,z,array[x,y,z]
        neighbors=0
        vol+=1

        for xDiff in [-1,0,1]:
            for yDiff in [-1,0,1]:
                for zDiff in [-1,0,1]:
                    if abs(xDiff)+abs(yDiff)+abs(zDiff)==1:
                        x1=x+xDiff
                        y1=y+yDiff
                        z1=z+zDiff
                        if x1>=0 and y1>=0 and z1>=0 and x1<xSize and y1<ySize and z1<zSize:
                            #print '-',x1,y1,z1,array[x1,y1,z1]
                            if array[x1,y1,z1]:
                                #print '-',x1,y1,z1,array[x1,y1,z1]
                                neighbors+=1

        #print 'had this many neighbors',neighbors
        sa+=(6-neighbors)

    onethird=float(1)/float(3)
    twothirds=float(2)/float(3)
    sph = ((np.pi**onethird)*((6*vol)**twothirds)) / sa
    #print 'sphericity',sphericity
    return sph

分析测试

#Imports
import pstats, cProfile
import numpy as np
import pyximport
pyximport.install(setup_args={"script_args":["--compiler=mingw32"], "include_dirs":np.get_include()}, reload_support=True) #Generate cython version

#Create fake array to calc sphericity
fakeArray=np.zeros((512,512,512))
fakeArray[200:300,200:300,200:300]=1

#Profiling stuff
cProfile.runctx("sphericity(fakeArray)", globals(), locals(), "Profile.prof")
s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("time").print_stats()

分析输出

Mon Oct 06 11:49:57 2014    Profile.prof

         12 function calls in 4.373 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    3.045    3.045    4.373    4.373 <string>:1(<module>)
        1    1.025    1.025    1.025    1.025 {method 'nonzero' of 'numpy.ndarray' objects}
        2    0.302    0.151    0.302    0.151 {numpy.core.multiarray.array}
        1    0.001    0.001    1.328    1.328 numeric.py:731(argwhere)
        1    0.000    0.000    0.302    0.302 fromnumeric.py:492(transpose)
        1    0.000    0.000    0.302    0.302 fromnumeric.py:38(_wrapit)
        1    0.000    0.000    0.000    0.000 {method 'transpose' of 'numpy.ndarray' objects}
        1    0.000    0.000    0.302    0.302 numeric.py:392(asarray)
        1    0.000    0.000    0.000    0.000 numeric.py:462(asanyarray)
        1    0.000    0.000    0.000    0.000 {getattr}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

Answer 1

Jaime 可能给出了一个很好的答案，但我会评论改进 Cython 代码并添加性能比较。

首先，您应该使用“注释”功能，cython -a filename.pyx，这将生成一个 HTML 文件。在浏览器中加载它，它会用橙黄色突出显示“慢”行，这表示可以进行改进的地方。

Annotate 立即揭示了两件很容易修复的事情:

将成语转换成 cython 可以理解的东西

首先，这些线路很慢:

        for xDiff in [-1,0,1]:
            for yDiff in [-1,0,1]:
                for zDiff in [-1,0,1]:

这是因为 Cython 不知道如何将列表迭代转换为干净的 C 代码。它需要转化为 Cython 可以优化的等效代码，即“范围内”形式:

        for xDiff in range(-1, 2):
            for yDiff in range(-1, 2):
                for zDiff in range(-1, 2):

用于快速索引的类型数组

接下来是这条线很慢:

                            if array[x1,y1,z1]:

这是因为 array 没有被赋予类型。因为它使用 python 级索引而不是 c 级索引。为了补救这个你需要给数组一个类型，这可以通过这种方式完成:

def sphericity(np.ndarray[np.uint8_t, ndim=3] array):

假设数组的类型是 'uint8'，替换为合适的类型(注意:Cython 不支持 'np.bool' 类型，因此我使用 'uint8')

你也可以使用内存 View ，你不能在内存 View 上使用 numpy 函数，但你可以在数组上创建一个 View ，然后索引 View 而不是数组:

    cdef np.uint8_t array_view [:, :, :] = array
    ...
                                    if array_view[x1,y1,z1]:

内存 View 可能会稍微快一些，并且在数组(python 级别调用)和 View (c 级别调用)之间做了明确的划分。如果您不使用 numpy 函数，则可以毫无问题地使用内存 View 。

重写代码以避免多次遍历数组

剩下的就是计算 index 和 nonZeroCount 都很慢，这是出于各种原因但主要只与数据的绝对大小有关(本质上，迭代超过 512*512*512 个元素只需要时间!)一般来说，Numpy 可以做的任何事情，经过优化的 Cython 可以做得更快(通常快 2-10 倍)——numpy 只会为您节省大量重新发明轮子和大量打字的时间，并让您在更高层次上思考(如果您不是也是一个 c 程序员，你可能无法很好地优化 cython)。但在这种情况下很简单，您只需删除 index 和 nonZeroCount 以及所有相关代码，然后执行以下操作:

    for x in range(0, xSize):
        for y in range(0, ySize):
            for z in range(0, zSize):
                if array[x,y,z] == 0:
                    continue
                ... 

这是非常快的，因为 c(干净的 Cython 编译成完美的)每秒执行数十亿次操作没有问题。通过消除 index 和 nonZeroCount 步骤，您实际上已经在整个数组上节省了两次完整的迭代，即使在最高速度下，每次迭代也至少需要大约 0.1 秒。更重要的是 CPU 缓存，整个数组是 128mb，比 cpu 缓存大得多，所以一次完成所有操作可以更好地利用 cpu 缓存(如果数组完全适合 CPU，多次通过就没那么重要了缓存)。

优化版本

这是我优化后的完整代码:

#cython: boundscheck=False, nonecheck=False, wraparound=False
import numpy as np
cimport numpy as np

def sphericity2(np.uint8_t [:, :, :] array):

    #Pass an mask array (1's are marked, 0's ignored)
    cdef int xSize,ySize,zSize
    xSize=array.shape[0]
    ySize=array.shape[1]
    zSize=array.shape[2]

    cdef int sa,vol,x,y,z,neighbors,xDiff,yDiff,zDiff,x1,y1,z1
    cdef float onethird,twothirds,sp
    sa=vol=0 #keep running tally of volume and surface area

    for x in range(0, xSize):
        for y in range(0, ySize):
            for z in range(0, zSize):
                if array[x,y,z] == 0:
                    continue

                neighbors=0
                vol+=1

                for xDiff in range(-1, 2):
                    for yDiff in range(-1, 2):
                        for zDiff in range(-1, 2):
                            if abs(xDiff)+abs(yDiff)+abs(zDiff)==1:
                                x1=x+xDiff
                                y1=y+yDiff
                                z1=z+zDiff
                                if x1>=0 and y1>=0 and z1>=0 and x1<xSize and y1<ySize and z1<zSize:
                                    #print '-',x1,y1,z1,array[x1,y1,z1]
                                    if array[x1,y1,z1]:
                                        #print '-',x1,y1,z1,array[x1,y1,z1]
                                        neighbors+=1

                #print 'had this many neighbors',neighbors
                sa+=(6-neighbors)

    onethird=float(1)/float(3)
    twothirds=float(2)/float(3)
    sph = ((np.pi**onethird)*((6*vol)**twothirds)) / sa
    #print 'sphericity',sphericity
    return sph

球形执行时间比较:

Original         : 2.123sJaime's          : 1.819sOptimized Cython : 0.136s@ moarningsun    : 0.090s

在所有 Cython 解决方案中，运行速度提高了 15 倍以上，展开的内部循环(参见评论)运行速度提高了 23 倍以上。