在 Cython 中使用 NumPy 函数对数组元素进行最小二乘拟合

Using NumPy functions in Cython for least-squares fitting of array elements

我需要编写一个脚本来对 4 张相似的 500x500 图像进行逐个像素的最小二乘拟合。比如,我需要将所有四个图像上特定像素位置的值拟合到一个长度为三的向量,对每个像素使用相同的 4x3 矩阵。

如果不对每个像素进行嵌套 for 循环迭代,我看不出有什么方法可以做到这一点,所以我认为 cython 可以加快速度。我以前从未使用过 cython,但我根据文档示例编写了以下代码。

问题是,这 运行 比纯 python 实施(~25 秒)慢或慢(~27 秒)。

有人看到是什么导致速度变慢了吗?谢谢!

import numpy as np
cimport numpy as np
cimport cython

npint = np.int16
npfloat = np.float64

ctypedef np.int16_t npint_t
ctypedef np.float64_t npfloat_t


@cython.boundscheck(False)
@cython.wraparound(False)

def fourbythree(np.ndarray[npfloat_t, ndim=2] U_mat, np.ndarray[npint_t, ndim=3] G):
    assert U_mat.dtype == npfloat and G.dtype == npint
    cdef unsigned int z = G.shape[0]
    cdef unsigned int rows = G.shape[1]
    cdef unsigned int cols = G.shape[2]
    cdef np.ndarray[npfloat_t, ndim= 3] a  = np.empty((z - 1, rows, cols), dtype=npfloat)
    cdef npfloat_t resid
    cdef unsigned int rank
    cdef Py_ssize_t row, col
    cdef np.ndarray s

    for row in range(rows):
        for col in range(cols):
            a[:, row, col] = np.linalg.lstsq(U_mat, G[:, row, col])[0]
    return a

您不需要迭代 - 您只需调用 lstsq 即可完成所有操作。 lstsq 允许第二个参数是二维的,在这种情况下结果也是二维的。您的数组是 3D 的,但是您可以轻松地将其重新整形为 2D,然后重新整形输出(并且重新整形基本上是免费的 - 它不需要复制数据):

a = np.linalg.lstsq(U_mat, G.reshape((G.shape[0],-1)))[0]
a = a.reshape((a.shape[0],G.shape[1],G.shape[2]))

这都是无类型的纯 Python 代码,因为这实际上不是任何索引,所以我不希望 Cython 提供帮助。

我从中得到了大约 400 倍的加速(尽管其中一些是因为 "one call" 版本似乎与 运行 并行出现,而 Cython 版本没有)。我认为加速的主要原因是重复调用 Python 函数的开销(假设它在非常小的数组上工作)。