优化通过每个元素迭代的NumPy矩阵之和

Question

我正在使用numpy 1.9,python 2.7和opencv,处理大矩阵,我必须多次进行以下操作

def sumShifted(A):  # A: numpy array 1000*1000*10
    return A[:, 0:-1] + A[:, 1:]    

如果可能的话,我想优化这个操作; 我尝试使用Cython,但我没有得到任何明显的改进,但我不排除它是因为我的实施不好.

有没有办法让它更快？

编辑: sumShifted在这样的for循环中调用:

for i in xrange(0, 400):
    # ... Various operations on B
    A = sumShifted(B)
    # ... Other operations on B


#More detailed
for i in xrange(0, 400):
    A = sumShifted(a11)
    B = sumShifted(a12)
    C = sumShifted(b12)
    D = sumShifted(b22)

    v = -upQ12/upQ11

    W, X, Z = self.function1( input_matrix, v, A, C[:,:,4], D[:,:,4] )
    S, D, F = self.function2( input_matrix, v, A, C[:,:,5], D[:,:,5] )
    AA      = self.function3( input_matrix, v, A, C[:,:,6], D[:,:,6] )
    BB      = self.function4( input_matrix, v, A, C[:,:,7], D[:,:,7] )

EDIT2:根据您的建议,我创建了这两个可运行的基准测试(使用Cython),将4 sumShifted种方法合并为一种.

A, B, C, D= improvedSumShifted(E, F, G, H)
#E,F: 1000x1000 matrices
#G,H: 1000x1000x8 matrices

#first implementation
def improvedSumShifted(np.ndarray[dtype_t, ndim=2] a, np.ndarray[dtype_t, ndim=2] b, np.ndarray[dtype_t, ndim=3] c, np.ndarray[dtype_t, ndim=3] d):
  cdef unsigned int i,j,k;
  cdef unsigned int w = a.shape[0], h = a.shape[1]-1, z = c.shape[2]
  cdef np.ndarray[dtype_t, ndim=2] aa = np.empty((w, h))
  cdef np.ndarray[dtype_t, ndim=2] bb = np.empty((w, h))
  cdef np.ndarray[dtype_t, ndim=3] cc = np.empty((w, h, z))
  cdef np.ndarray[dtype_t, ndim=3] dd = np.empty((w, h, z))
  with cython.boundscheck(False), cython.wraparound(False), cython.overflowcheck(False), cython.nonecheck(False):
    for i in range(w):
      for j in range(h):
        aa[i,j] = a[i,j] + a[i,j+1]
        bb[i,j] = b[i,j] + b[i,j+1]
        for k in range(z):
          cc[i,j,k] = c[i,j,k] + c[i,j+1,k]
          dd[i,j,k] = d[i,j,k] + d[i,j+1,k]
return aa, bb, cc, dd

#second implementation
def improvedSumShifted(np.ndarray[dtype_t, ndim=2] a, np.ndarray[dtype_t, ndim=2] b, np.ndarray[dtype_t, ndim=3] c, np.ndarray[dtype_t, ndim=3] d):
  cdef unsigned int i,j,k;
  cdef unsigned int w = a.shape[0], h = a.shape[1]-1, z = c.shape[2]
  cdef np.ndarray[dtype_t, ndim=2] aa = np.copy(a[:, 0:h])
  cdef np.ndarray[dtype_t, ndim=2] bb = np.copy(b[:, 0:h])
  cdef np.ndarray[dtype_t, ndim=3] cc = np.copy(c[:, 0:h])
  cdef np.ndarray[dtype_t, ndim=3] dd = np.copy(d[:, 0:h])
  with cython.boundscheck(False), cython.wraparound(False), cython.overflowcheck(False), cython.nonecheck(False):
  for i in range(w):
    for j in range(h):
      aa[i,j] += a[i,j+1]
      bb[i,j] += b[i,j+1]
      for k in range(z):
        cc[i,j,k] += c[i,j+1,k]
        dd[i,j,k] += d[i,j+1,k]

return aa, bb, cc, dd

Answer 1

这个函数不太可能进一步加速:它在python级别上只进行了四次操作:

1. (2x)对输入执行切片.这些切片非常快,因为它们只需要一些整数运算来计算新的步幅和大小.
2. 为输出分配新数组.对于这样一个简单的功能,这是一个很大的负担.
3. 评估np.add两个切片上的ufunc,这是一个在numpy中高度优化的操作.

实际上,我的基准测试显示无论是使用numba还是使用cython都没有改进.在我的机器上,如果输出数组是预先分配的,我每次调用的持续时间约为30毫秒,如果考虑内存分配,则为50毫秒.

纯粹的numpy版本:

import numpy as np

def ss1(A):
    return np.add(A[:,:-1,:],A[:,1:,:])

def ss2(A,output):
    return np.add(A[:,:-1,:],A[:,1:,:],output)

cython版本:

import numpy as np
cimport numpy as np
cimport cython

def ss3(np.float64_t[:,:,::1] A not None):
    cdef unsigned int i,j,k;
    cdef np.float64_t[:,:,::1] ret = np.empty((A.shape[0],A.shape[1]-1,A.shape[2]),'f8')
    with cython.boundscheck(False), cython.wraparound(False):
        for i in range(A.shape[0]):
            for j in range(A.shape[1]-1):
                for k in range(A.shape[2]):
                    ret[i,j,k] = A[i,j,k] + A[i,j+1,k]
    return ret

def ss4(np.float64_t[:,:,::1] A not None, np.float64_t[:,:,::1] ret not None):
    cdef unsigned int i,j,k;
    assert ret.shape[0]>=A.shape[0] and ret.shape[1]>=A.shape[1]-1 and ret.shape[2]>=A.shape[2]
    with cython.boundscheck(False), cython.wraparound(False):
        for i in range(A.shape[0]):
            for j in range(A.shape[1]-1):
                for k in range(A.shape[2]):
                    ret[i,j,k] = A[i,j,k] + A[i,j+1,k]
    return ret

numba版本(当前numba 0.14.0无法在优化函数中分配新数组):

@numba.njit('f8[:,:,:](f8[:,:,:],f8[:,:,:])')
def ss5(A,output):
    for i in range(A.shape[0]):
        for j in range(A.shape[1]-1):
            for k in range(A.shape[2]):
                output[i,j,k] = A[i,j,k] + A[i,j+1,k]
    return output

以下是时间安排:

>>> A = np.random.randn((1000,1000,10))
>>> output = np.empty((A.shape[0],A.shape[1]-1,A.shape[2]))

>>> %timeit ss1(A)
10 loops, best of 3: 50.2 ms per loop

>>> %timeit ss2(A,output)
10 loops, best of 3: 30.8 ms per loop

>>> %timeit ss3(A)
10 loops, best of 3: 50.8 ms per loop

>>> %timeit ss4(A,output)
10 loops, best of 3: 30.9 ms per loop

>>> %timeit ss5(A,output)
10 loops, best of 3: 31 ms per loop