如何优化计算三维数组沿 Z 轴的平均值的速度？Cython vs Numpy

Question

我试图加速对3D数组沿Z轴的平均值计算。我阅读了Cython文档以添加类型、内存视图等来完成此任务。然而，当我比较基于NumPy和基于Cython语法及编译生成.so文件的函数时，前者胜过后者。请问我的代码中是否有哪些步骤或类型声明存在问题或遗漏？

这是我的NumPy版本：python_mean.py

    import numpy as np


    def mean_py(array):
        x = array.shape[1]
        y = array.shape[2]
        values = []
        for i in range(x):
            for j in range(y):
                values.append((np.mean(array[:, i, j])))

        values = np.array([values])
        values = values.reshape(500,500)
        return values

这是我的 cython_mean.pyx 文件。

     %%cython
     from cython import wraparound, boundscheck
     import numpy as np
     cimport numpy as np 

     DTYPE = np.double

     @boundscheck(False)
     @wraparound(False)
     def cy_mean(double[:,:,:] array):
        cdef Py_ssize_t x_max = array.shape[1]
        cdef Py_ssize_t y_max = array.shape[2]
        cdef double[:,:] result = np.zeros([x_max, y_max], dtype = DTYPE)
        cdef double[:,:] result_view = result
        cdef Py_ssize_t i,j
        cdef double mean
        cdef list values 
        for i in range(x_max):
            for j in range(y_max):
                mean = np.mean(array[:,i,j])
                result_view[i,j] = mean
        return result

当我导入两个函数并开始对一个3D numpy数组进行计算时，我得到了以下结果：

    import numpy as np
    a = np.random.randn(250_000)
    b = np.random.randn(250_000)
    c = np.random.randn(250_000)

    array = np.vstack((a,b,c)).reshape(3, 500, 500)

    import mean_py
    from mean_py import mean_py
    %timeit mean_py(array)


    4.82 s ± 84.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)



    import cython_mean
    from cython_mean import cy_mean


    7.3 s ± 499 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

为什么Cython代码性能如此低？感谢您的帮助。

- Roger Almengor

---

Cython是优化Python代码的通用解决方案。NumPy是数学计算的特定解决方案。因此，在数学计算方面，大多数情况下应该选择NumPy... - Laurent LAPORTE

始终使用%%cython -a查看实际情况。问题在于使用np.mean()。如果您在循环中写出np.mean，您可以轻松地达到Numpy的性能（很可能numpy实现也是用Cython编写的）。 - max9111

Answer 1

Numpy解决方案

对于这个特定的问题，使用numpy.mean的axis参数可能是最快的实现方式（即values = np.mean(array, axis=0)）。

如下所示的基准测试表明，使用你的示例，numpy.mean的速度几乎快了1000倍。

In []: %timeit mean_py(array)
1.23 s ± 3.99 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In []: %timeit array.mean(0)
1.07 ms ± 3.76 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In []: np.all(array.mean(0) == mean_py(array))
Out[]: True

---

你原始方法的建议

这不是为什么 cython 版本没有更快的解释，而是为了改进仅使用 numpy 的版本（避免将 list 作为（慢速）中间数据结构）的建议：

    import numpy as np


    def mean_py(array):
        x = array.shape[1]
        y = array.shape[2]
        #avoid creating values as list first
        #and create empty array instead
        values = np.empty((x,y), type(array[0][0][0]))
        for i in range(x):
            for j in range(y):
                #no more need for append operations
                values[i][j] = ((np.mean(array[:, i, j])))

        #no more need for conversion to array
        #values = np.array([values])
        #values = values.reshape(500,500)
        return values

- Pete

---

很遗憾，当我试图将计算出的平均值分配给数组中的位置时，收到了以下错误消息：IndexError: index 500 is out of bounds for axis 0 with size 500。 - Roger Almengor