Cython，numpy 加速

Question

我正在尝试编写一种算法来计算二维数组的某些相邻元素的平均值。

想看看能不能用Cython提速，我自己还是第一次用

Python 版本:

import numpy as np

def clamp(val, minval, maxval):
    return max(minval, min(val, maxval))


def filter(arr, r):
    M = arr.shape[0]
    N = arr.shape[1]

    new_arr = np.zeros([M, N], dtype=np.int)

    for x in range(M):
        for y in range(N):
            # Corner elements
            p1 = clamp(x-r, 0, M)
            p2 = clamp(y-r, 0, N)
            p3 = clamp(y+r, 0, N-1)
            p4 = clamp(x+r, 0, M-1)

            nbr_elements = (p3-p2-1)*2+(p4-p1-1)*2+4

            tmp = 0

            # End points
            tmp += arr[p1, p2]
            tmp += arr[p1, p3]
            tmp += arr[p4, p2]
            tmp += arr[p4, p3]

            # The rest
            tmp += sum(arr[p1+1:p4, p2])
            tmp += sum(arr[p1+1:p4, p3])
            tmp += sum(arr[p1, p2+1:p3])
            tmp += sum(arr[p4, p2+1:p3])

            new_arr[x, y] = tmp/nbr_elements

    return new_arr

以及我对 Cython 实现的尝试。我发现如果重新实现 max/min/sum 比使用 python 版本

更快

Cython 版本:

from __future__ import division
import numpy as np
cimport numpy as np

DTYPE = np.int
ctypedef np.int_t DTYPE_t

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b

def clamp(int val, int minval, int maxval):
    return int_max(minval, int_min(val, maxval))

def cython_sum(np.ndarray[DTYPE_t, ndim=1] y):
    cdef int N = y.shape[0]
    cdef int x = y[0]
    cdef int i
    for i in xrange(1, N):
        x += y[i]
    return x


def filter(np.ndarray[DTYPE_t, ndim=2] arr, int r):
    cdef M = im.shape[0]
    cdef N = im.shape[1]

    cdef np.ndarray[DTYPE_t, ndim=2] new_arr = np.zeros([M, N], dtype=DTYPE)
    cdef int p1, p2, p3, p4, nbr_elements, tmp

    for x in range(M):
        for y in range(N):
            # Corner elements
            p1 = clamp(x-r, 0, M)
            p2 = clamp(y-r, 0, N)
            p3 = clamp(y+r, 0, N-1)
            p4 = clamp(x+r, 0, M-1)

            nbr_elements = (p3-p2-1)*2+(p4-p1-1)*2+4

            tmp = 0

            # End points
            tmp += arr[p1, p2]
            tmp += arr[p1, p3]
            tmp += arr[p4, p2]
            tmp += arr[p4, p3]

            # The rest
            tmp += cython_sum(arr[p1+1:p4, p2])
            tmp += cython_sum(arr[p1+1:p4, p3])
            tmp += cython_sum(arr[p1, p2+1:p3])
            tmp += cython_sum(arr[p4, p2+1:p3])

            new_arr[x, y] = tmp/nbr_elements

    return new_arr

我做了一个测试脚本：

import time
import numpy as np

import square_mean_py
import square_mean_cy

N = 500

arr = np.random.randint(15, size=(N, N))
r = 8

# Timing

t = time.time()
res_py = square_mean_py.filter(arr, r)
print time.time()-t

t = time.time()
res_cy = square_mean_cy.filter(arr, r)
print time.time()-t

打印

9.61458301544
1.44476890564

这是大约的加速。 7次。我已经看到很多 Cython implmentations 产生了更好的加速，所以我在想也许你们中的一些人看到了加速算法的潜在方法？

Answer 1

您的 Cython 脚本存在一些问题：

您没有向 Cython 提供一些关键信息，例如范围中使用的 x, y, M 和 N 的类型。
我已经 cdef 编辑了 cython_sum 和 clamp 这两个函数，因为您在 Python 级别不需要它们。
filter函数中出现的im是什么？我假设你的意思是 arr.

修复这些问题我会 rewrite/modify 你的 Cython 脚本像这样：

from __future__ import division
import numpy as np
cimport numpy as np
from cython cimport boundscheck, wraparound 

DTYPE = np.int
ctypedef np.int_t DTYPE_t

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b

cdef int clamp3(int val, int minval, int maxval):
    return int_max(minval, int_min(val, maxval))

@boundscheck(False)
cdef int cython_sum2(DTYPE_t[:] y):
    cdef int N = y.shape[0]
    cdef int x = y[0]
    cdef int i
    for i in range(1, N):
        x += y[i]
    return x

@boundscheck(False)
@wraparound(False)
def filter3(DTYPE_t[:,::1] arr, int r):
    cdef int M = arr.shape[0]
    cdef int N = arr.shape[1]

    cdef np.ndarray[DTYPE_t, ndim=2, mode='c'] \
    new_arr = np.zeros([M, N], dtype=DTYPE)
    cdef int p1, p2, p3, p4, nbr_elements, tmp, x, y

    for x in range(M):
        for y in range(N):
            # Corner elements
            p1 = clamp3(x-r, 0, M)
            p2 = clamp3(y-r, 0, N)
            p3 = clamp3(y+r, 0, N-1)
            p4 = clamp3(x+r, 0, M-1)

            nbr_elements = (p3-p2-1)*2+(p4-p1-1)*2+4

            tmp = 0

            # End points
            tmp += arr[p1, p2]
            tmp += arr[p1, p3]
            tmp += arr[p4, p2]
            tmp += arr[p4, p3]

            # The rest
            tmp += cython_sum2(arr[p1+1:p4, p2])
            tmp += cython_sum2(arr[p1+1:p4, p3])
            tmp += cython_sum2(arr[p1, p2+1:p3])
            tmp += cython_sum2(arr[p4, p2+1:p3])

            new_arr[x, y] = <int>(tmp/nbr_elements)

    return new_arr

这是我机器上的时间：

arr = np.random.randint(15, size=(500, 500))

Original (Python) version: 7.34 s
Your Cython version: 1.98 s
New Cython version: 0.0323 s

这比您的 Cython 脚本快了将近 60 倍，比原始 Python 脚本快了 200 多倍。

Cython，numpy 加速

Cython, numpy speed-up

python

performance

numpy

cython