在numpy的示例代碼中正確使用Cython

Question

我想知道在使用Cython和Numpy時我是否遺漏了一些東西，因為我沒有看到太多改進。 我把這段代碼寫成了一個例子。

朴素版：

import numpy as np
from skimage.util import view_as_windows

it = 16
arr = np.arange(1000*1000, dtype=np.float64).reshape(1000,1000)
windows = view_as_windows(arr, (it, it), it)
container = np.zeros((windows.shape[0], windows.shape[1]))
def test(windows):
    for i in range(windows.shape[0]):
        for j in range(windows.shape[1]):
            container[i,j] = np.mean(windows[i,j])
    return container

%%timeit 

test(windows)
1 loops, best of 3: 131 ms per loop

Cythonized版本：

%%cython --annotate

import numpy as np
cimport numpy as np
from skimage.util import view_as_windows
import cython
cdef int step = 16

arr = np.arange(1000*1000, dtype=np.float64).reshape(1000,1000)
windows = view_as_windows(arr, (step, step), step)

@cython.boundscheck(False)
def cython_test(np.ndarray[np.float64_t, ndim=4]  windows):
    cdef np.ndarray[np.float64_t, ndim=2] container = np.zeros((windows.shape[0], windows.shape[1]),dtype=np.float64)
    cdef int i, j
    I = windows.shape[0]
    J = windows.shape[1]
    for i in range(I):
        for j in range(J):
            container[i,j] = np.mean(windows[i,j])
    return container


%timeit cython_test(windows)
10 loops, best of 3: 126 ms per loop

正如你所看到的，有一個非常適度的改進，所以也許我做錯了。 順便說一下，Cython產生的注釋如下：

正如您所看到的，即使在包含有效的索引語法np.ndarray[DTYPE_t, ndim=2]之后，numpy線也具有黃色背景。 為什么？

順便說一句，在我看來，理想的結果是能夠使用大多數numpy函數，但在利用有效的索引語法或者像HYRY的答案中的內存視圖之后仍然得到一些合理的改進。

UPDATE

似乎我在上面發布的代碼中沒有做任何錯誤，並且某些行中的黃色背景是正常的，所以我想知道以下情況：在哪些情況下我可以從鍵入cdef np.ndarray[np.float64_t, ndim=2]獲益numpy數組前面的cdef np.ndarray[np.float64_t, ndim=2] ？ 我想有一些具體的例子，這有用，否則就沒有太多的目的。

Answer 1

你需要自己實現mean()函數來加速代碼，這是因為調用numpy函數的開銷非常高。

@cython.boundscheck(False)
@cython.wraparound(False)
def cython_test(double[:, :, :, :]  windows):
    cdef double[:, ::1] container
    cdef int i, j, k, l
    cdef int n0, n1, n2, n3
    cdef double inv_n
    cdef double s
    n0, n1, n2, n3 = windows.base.shape
    container = np.zeros((n0, n1))
    inv_n = 1.0 / (n2 * n3)
    for i in range(n0):
        for j in range(n1):
            s = 0
            for k in range(n2):
                for l in range(n3):
                    s += windows[i, j, k, l]
            container[i,j] = s * inv_n
    return container.base

這是%timeit結果：

• python_test(windows) ：63.7 ms
• cython_test(windows) ：1.24毫秒
• np.mean(windows, axis=(2, 3)) ：2.66 ms