在遍历数组时，Cupy 比 numpy 慢

Question

我有代码，我想与cupy并行化。我认为这将是直截了当的 - 只需写"import cupy as cp"，并将我写的 np. 替换为 cp.，它就会起作用。

而且，它确实有效，代码确实运行，但需要慢得多。我认为在遍历更大的数组时，与 numpy 相比，它最终会更快，但似乎从未发生过。

代码为：

q = np.zeros((5,5))
q[:,0] = 20
def foo(array):
    result = array
    shedding_row = array*0
    for i in range((array.shape[0])):
        for j in range((array.shape[1])-1):
            shedding_param = 2 * (result[i,j])**.5             
            shedding = (np.random.poisson( (shedding_param), 1))[0]
            if shedding >= result[i,j]:
                shedding = result[i,j] - 1
            result[i,j+1] = result[i,j] - shedding
            if result[i,j+1]<0:
                result[i,j+1] = 0
            shedding_row[i,j+1] = shedding  
    return(result,shedding_row)
x,y = foo(q)

这应该用cupy变得更快吗？我用错了吗？

Answer 1

为了获得numpy或cupy的快速性能，你应该使用并行操作而不是使用 for 循环。

举个例子，

for i in range((array.shape[0])):
    for j in range((array.shape[1])-1):
        shedding_param = 2 * (result[i,j])**.5

这可以计算为

xp = numpy  # change to cupy for GPU
shedding_param = 2 * xp.sqrt(result[:, :-1])