python - 给数组切片赋值很慢

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. 在条件 block 之后，将 v1、v2、v3 分配给 res[i,j,0] 等三个单独的分配。