为数组切片分配值很慢

Question

我试图通过在Cython中实现它来优化Python算法.我的问题是关于以下代码中存在的某个性能瓶颈:

@cython.boundscheck(False) # turn off bounds-checking for entire function
def anglesToRGB( np.ndarray[double, ndim=2] y, np.ndarray[double, ndim=2] x ):

cdef double angle
cdef double Hp
cdef double C
cdef double X
cdef np.ndarray[double, ndim=3] res = np.zeros([y.shape[0], y.shape[1], 3], dtype=np.float64)

for i in xrange(y.shape[0]):
    for j in xrange(y.shape[1]):
        angle = atan2( y[i,j], x[i,j] )*180.0/PI+180

        C = sqrt(pow(y[i,j],2)+pow(x[i,j],2))/360.0 #Chroma
        Hp = angle/60.0
        X = C*(1-fabs( Hp%2-1))

        C *= 255
        X *= 255

        if (0. <= Hp < 1.):
            res[i,j,:] = [C,X,0]
        elif (1. <= Hp < 2.):
            res[i,j,:] = [X,C,0]
        elif (2. <= Hp < 3.):
            res[i,j,:] = [0,C,X]
        elif (3. <= Hp < 4.):
            res[i,j,:] = [0,X,C]
        elif (4. <= Hp < 5.):
            res[i,j,:] = [X,C,C]
        else:
            res[i,j,:] = [C,0,X]

return res

我已经确定了当我为res数组的一个片段分配值列表时的主要瓶颈,例如

res[i,j,:] = [C,X,0]

但是,如果我将作业更改为

res[i,j,0] = C
res[i,j,1] = X
res[i,j,2] = 0

然后代码运行速度提高了几个数量级.对我来说这很奇怪,因为Cython编译器必须足够聪明才能为我做到这一点？或者我是否需要首先提供一些提示？我应该注意到,将切片更改为0:3而不是:并使值列表成为一个numpy数组并不会提高性能.

我想知道的是为什么这个操作如此糟糕地杀死性能,如果有任何方法可以解决它而不必牺牲方便的列表和切片符号.

最好的祝福

Answer 1

不，Cython（使用 0.17 进行测试）不够智能，无法优化此切片分配。如果您查看生成的 C 代码（使用cython -a并单击 HTML 报告中的任意行来查看生成的代码），那么您可以看到

res[i,j,:] = [C,X,0]

被编译为

• C 和 Python 浮点类型之间的转换
• 列表的分配[C,X,0]
• 元组的分配(i, j, slice(None))
• 打电话给res.__setitem__
• 对所有这些进行错误检查
• 已分配结构的解除分配

即，执行此代码时 CPython 会执行几乎所有相同的操作。

为了解决这个问题，你可以采取以下措施：

1. 声明三个变量，例如cdef double v1, v2, v3；
2. 在条件语句中分配给这些，例如v1, v2, v3 = C, X, 0等，它被优化为三个 C 分配；
3. 在条件块之后，在三个单独的赋值中分配给等v1, v2, v3。res[i,j,0]