意外缓慢的cython卷积代码

Question

我需要快速计算一个矩阵，这个矩阵的每个元素是通过将一个滤波器和一个向量进行卷积得到的。然后我会对得到的向量进行下采样，最后再和另一个向量做点积。具体来说，我想计算

M = [conv(e_j, f)*P_i*v_i ]_{i,j},

其中 i 从 1 到 n，j 从 1 到 m。这里的 e_j 是一个大小为 n 的指示向量，只有在第 j 列有一个 1，f 是长度为 s 的滤波器，P_i 是一个 (n+s-1)-行 k 列的矩阵，用来从卷积结果中取出合适的 k 个元素，而 v_i 是一个长度为 k 的列向量。

计算 M 的每个元素需要 O(n*s) 次操作，所以计算整个 M 需要 O(n*s*n*m) 次操作。比如当 n=6，m=7，s=3 时，我电脑的一个核心（8GLOPs）应该能在大约 0.094 微秒内计算出 M。然而，我的一个非常简单的 Cython 实现，按照 Cython 文档中的示例，却需要超过 2 毫秒来计算这些参数的示例。这差不多有 4 个数量级的差别！

这里有一个包含 Cython 实现和测试代码的 shar 文件。你可以把它复制粘贴到一个文件中，然后在一个干净的目录下运行 'bash <fname>' 来获取代码，然后运行 'bash ./test.sh' 来看看性能有多糟糕。

cat > fastcalcM.pyx <<'EOF'

import numpy as np
cimport numpy as np
cimport cython
from scipy.signal import convolve

DTYPE=np.float32
ctypedef np.float32_t DTYPE_t

@cython.boundscheck(False)
def calcM(np.ndarray[DTYPE_t, ndim=1, negative_indices=False] filtertaps, int
        n, int m, np.ndarray[np.int_t, ndim=2, negative_indices=False]
        keep_indices, np.ndarray[DTYPE_t, ndim=2, negative_indices=False] V):
    """ Computes a numrows-by-k matrix M whose entries satisfy
        M_{i,k} = [conv(e_j, f)^T * P_i * v_i],
        where v_i^T is the i-th row of V, and P_i samples the entries from
        conv(e_j, f)^T indicated by the ith row of the keep_indices matrix """

    cdef int k = keep_indices.shape[1]

    cdef np.ndarray M = np.zeros((n, m), dtype=DTYPE)
    cdef np.ndarray ej = np.zeros((m,), dtype=DTYPE)
    cdef np.ndarray convolution
    cdef int rowidx, colidx, kidx

    for rowidx in range(n):
        for colidx in range(m):
            ej[colidx] = 1
            convolution = convolve(ej, filtertaps, mode='full')
            for kidx in range(k):
                M[rowidx, colidx] += convolution[keep_indices[rowidx, kidx]] * V[rowidx, kidx]
            ej[colidx] = 0

    return M

EOF
#-----------------------------------------------------------------------------
cat > test_calcM.py << 'EOF'

import numpy as np
from fastcalcM import calcM

filtertaps = np.array([-1, 2, -1]).astype(np.float32)
n, m = 6, 7
keep_indices = np.array([[1, 3], 
                         [4, 5],
                         [2, 2], 
                         [5, 5], 
                         [3, 4], 
                         [4, 5]]).astype(np.int)
V = np.random.random_integers(-5, 5, size=(6, 2)).astype(np.float32)

print calcM(filtertaps, n, m, keep_indices, V)

EOF
#-----------------------------------------------------------------------------
cat > test.sh << 'EOF'

python setup.py build_ext --inplace
echo -e "%run test_calcM\n%timeit calcM(filtertaps, n, m, keep_indices, V)" > script.ipy
ipython script.ipy

EOF
#-----------------------------------------------------------------------------
cat > setup.py << 'EOF'

from distutils.core import setup
from Cython.Build import cythonize
import numpy

setup(
    name="Fast convolutions",
    include_dirs = [numpy.get_include()],
    ext_modules = cythonize("fastcalcM.pyx")
)

EOF

我想也许是调用 scipy 的卷积函数导致了这个问题（我不确定 Cython 和 scipy 是否兼容），所以我按照 Cython 文档中的示例实现了自己的卷积代码，但结果是整体代码变得慢了大约 10 倍。

有没有什么想法可以让我更接近理论上可能的速度，或者为什么差距会这么大？

性能优化 scipy 数值计算 cython 算法复杂度卷积矩阵计算下采样

意外缓慢的cython卷积代码

1 个回答

撰写回答