python - numpy 函数 cythonization

Question

我在纯 python 中有以下函数:

import numpy as np

def subtractPython(a, b):
    xAxisCount = a.shape[0]
    yAxisCount = a.shape[1]

    shape = (xAxisCount, yAxisCount, xAxisCount)
    results = np.zeros(shape)
    for index in range(len(b)):
        subtracted = (a - b[index])
        results[:, :, index] = subtracted
    return results

我试着用这种方式对它进行 cythonize:

import numpy as np
cimport numpy as np

DTYPE = np.int
ctypedef np.int_t DTYPE_t

def subtractPython(np.ndarray[DTYPE_t, ndim=2] a, np.ndarray[DTYPE_t, ndim=2] b):
    cdef int xAxisCount = a.shape[0]
    cdef int yAxisCount = a.shape[1]

    cdef np.ndarray[DTYPE_t, ndim=3] results = np.zeros([xAxisCount, yAxisCount, xAxisCount], dtype=DTYPE)

    cdef int lenB = len(b)

    cdef np.ndarray[DTYPE_t, ndim=2] subtracted
    for index in range(lenB):
        subtracted = (a - b[index])
        results[:, :, index] = subtracted
    return results

但是，我没有看到任何加速。我是否遗漏了什么或无法加快此过程？

编辑 -> 我意识到我实际上并没有在上面的代码中对减法算法进行 cythonizing。我已经设法对其进行了 cythonize，但它具有与 a - b[:, None] 完全相同的运行时间，所以我猜这是该操作的最大速度。

这基本上是 a - b[:, None] -> 具有相同的运行时间

%%cython

import numpy as np
cimport numpy as np


DTYPE = np.int
ctypedef np.int_t DTYPE_t

cimport cython
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function
def subtract(np.ndarray[DTYPE_t, ndim=2] a, np.ndarray[DTYPE_t, ndim=2] b):
    cdef np.ndarray[DTYPE_t, ndim=3] result = np.zeros([b.shape[0], a.shape[0], a.shape[1]], dtype=DTYPE)

    cdef int lenB = b.shape[0]
    cdef int lenA = a.shape[0]
    cdef int lenColB = b.shape[1]

    cdef int rowA, rowB, column

    for rowB in range(lenB):
        for rowA in range(lenA):
            for column in range(lenColB):
                result[rowB, rowA, column] = a[rowA, column] - b[rowB, column]
    return result

Answer 1

当尝试优化一个功能时，人们总是应该知道这个功能的瓶颈是什么——否则你会花很多时间在错误的方向上运行。

让我们使用你的 python 函数作为基线(实际上我使用 result=np.zeros(shape,dtype=a.dtype) 否则你的方法返回 floats 这是可能是一个错误):

>>> import numpy as np
>>> a=np.random.randint(1,1000,(300,300), dtype=np.int)
>>> b=np.random.randint(1,1000,(300,300), dtype=np.int)
>>> %timeit subtractPython(a,b)
274 ms ± 3.61 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我们应该问自己的第一个问题是:此任务受内存或 CPU 限制吗？显然，这是一项受内存限制的任务 - 与所需的内存读取和写入访问相比，减法算不了什么。

这意味着，以上我们必须优化内存布局以减少缓存未命中。根据经验，我们的内存访问应该访问一个接一个的连续内存地址。

是这样吗？不，数组 result 是 C 顺序的，即行优先顺序，因此访问

results[:, :, index] = subtracted

不是连续的。另一方面，

results[index, :, :] = subtracted

将是连续访问。让我们改变信息在 result 中的存储方式:

def subtract1(a, b):
    xAxisCount = a.shape[0]
    yAxisCount = a.shape[1]

    shape = (xAxisCount,  xAxisCount, yAxisCount) #<=== Change order
    results = np.zeros(shape, dtype=a.dtype)
    for index in range(len(b)):
        subtracted = (a - b[index])
        results[index, :, :] = subtracted   #<===== consecutive access
    return results

现在的时间是:

>>> %timeit subtract1(a,b)
>>> 35.8 ms ± 285 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

还有 2 个小改进:我们不必用零初始化结果，我们可以节省一些 python 开销，但这只给了我们大约 5%:

def subtract2(a, b):
    xAxisCount = a.shape[0]
    yAxisCount = a.shape[1]

    shape = (xAxisCount,  xAxisCount, yAxisCount) 
    results = np.empty(shape, dtype=a.dtype)        #<=== no need for zeros
    for index in range(len(b)):
        results[index, :, :] = (a-b[index])   #<===== less python overhead
    return results

>>> %timeit subtract2(a,b)
34.5 ms ± 203 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

现在这比原始版本快大约 8 倍。

您可以使用 Cython 尝试进一步加速 - 但任务可能仍然受内存限制，所以不要期望它会显着加快 - 毕竟 cython 无法使内存工作得更快。然而，如果没有适当的分析，很难说有多少改进是可能的——如果有人想出更快的版本，也不会感到惊讶。