python - cython numpy 累积函数

Question

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

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

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

Answer 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。