cython numpy 累积函数答案

【问题标题】：cython numpy accumulate functioncython numpy 累积函数
【发布时间】：2012-05-31 13:38:42
【问题描述】：

我需要实现一个函数，用于对具有可变节长度的数组元素求和。所以，

a = np.arange(10)
section_lengths = np.array([3, 2, 4])
out = accumulate(a, section_lengths)
print out
array([  3.,   7.,  35.])

我在这里尝试了cython 的实现：

https://gist.github.com/2784725

对于性能，我将在 section_lengths 都相同的情况下与纯 numpy 解决方案进行比较：

LEN = 10000
b = np.ones(LEN, dtype=np.int) * 2000
a = np.arange(np.sum(b), dtype=np.double)
out = np.zeros(LEN, dtype=np.double)

%timeit np.sum(a.reshape(-1,2000), axis=1)
10 loops, best of 3: 25.1 ms per loop

%timeit accumulate.accumulate(a, b, out)
10 loops, best of 3: 64.6 ms per loop

您对提高性能有什么建议吗？

【问题讨论】：

我实现了几个建议，在github上查看更新版本：gist.github.com/2784725/…，仍然需要63ms，所以没有明显改进
这可能不是重点，但我想我会提到...... numpy 已经为 all ufunc 提供了一些接近它的东西。 np.add.reduceat(a, section_lengths.cumsum())。它必须稍作更改（cumsum 在开始时缺少 0，而您会得到额外的结束切片）并且您可能可以使用 cython 超越速度，但它是一个非常好的功能/技巧。

标签： python performance numpy cython

【解决方案1】：

您可以尝试以下方法：

除了@cython.boundscheck(False)编译器指令，还可以尝试添加@cython.wraparound(False)
在您的 setup.py 脚本中，尝试添加一些优化标志：

ext_modules = [Extension("accumulate", ["accumulate.pyx"], extra_compile_args=["-O3",])]
查看cython -a accumulate.pyx 生成的 .html 文件，看看是否有部分缺少静态类型或严重依赖 Python C-API 调用：

http://docs.cython.org/src/quickstart/cythonize.html#determining-where-to-add-types
在方法末尾添加return 语句。目前，它在i_el += 1 的紧密循环中进行了一堆不必要的错误检查。
不确定它是否会有所作为，但我倾向于制作循环计数器 cdef unsigned int 而不仅仅是 int

当section_lengths 不相等时，您还可以将您的代码与 numpy 进行比较，因为它可能需要的不仅仅是简单的sum。

【讨论】：

谢谢！我实施了您的所有建议，但仍然没有显着改善。感谢您建议 cython -a，对此一无所知。我添加了一个 return 语句，它显示了代码正在执行的一些奇怪的检查，请参阅gist.github.com/2784725#gistcomment-330807
我接受这个答案是因为它提供了有用的建议，但没有一个能带来显着的改进。如果其他人发现更好的东西，我会更改已接受的答案。

【解决方案2】：

nest for 循环更新out[i_bas] 很慢，你可以创建一个临时变量来做累加，当nest for 循环完成时更新out[i_bas]。以下代码将与 numpy 版本一样快：

import numpy as np
cimport numpy as np

ctypedef np.int_t DTYPE_int_t
ctypedef np.double_t DTYPE_double_t

cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
def accumulate(
       np.ndarray[DTYPE_double_t, ndim=1] a not None,
       np.ndarray[DTYPE_int_t, ndim=1] section_lengths not None,
       np.ndarray[DTYPE_double_t, ndim=1] out not None,
       ):
    cdef int i_el, i_bas, sec_length, lenout
    cdef double tmp
    lenout = out.shape[0]
    i_el = 0
    for i_bas in range(lenout):
        tmp = 0
        for sec_length in range(section_lengths[i_bas]):
            tmp += a[i_el]
            i_el+=1
        out[i_bas] = tmp

【讨论】：

谢谢！听从了你的建议，但没有明显改善，我在 github 上更新了我的版本