python - 并行重新分配大型数组中的元素

Question

我有一个 numpy 数组，chop_preds，它非常大(约 1000 万个元素)，需要修改以使其包含 1.0、0.5 或 0 的值(见下文) ).

我怎样才能并行化这个重新分配？

chop_preds=chop_preds.flatten()

for k in range(len(chop_preds)):
    if(chop_preds[k]>=0.4):
        chop_preds[k]=1.0 
    elif(chop_preds[k]<0.1):
        chop_preds[k]=0 
    else:
        chop_preds[k]=0.5 

my_sum=np.sum(chop_preds)

Answer 1

如果 chop_preds 已经是一个 numpy 数组，您可以使用:

chop_preds_flat = chop_preds.flatten()
chop_preds = 0.5 * np.ones_like(chop_preds_flat)
chop_preds[chop_preds_flat >= 0.4] = 1.
chop_preds[chop_preds_flat < 0.1] = 0.

my_sum = chop_preds.sum()

或者，如果您真的只需要总和，请使用 numpy.count_nonzero在这些选择上:

my_sum = 0.5 * np.count_nonzero((chop_preds_flat >= 0.1) & (chop_preds_flat < 0.4))
my_sum += np.count_nonzero(chop_preds_flat >= 0.4)

更简单，但更难阅读:

my_sum = ((chop_preds_flat >= 0.4) + 0.5 * ((chop_preds_flat >= 0.1) & (chop_preds_flat < 0.4))).sum()

在这三种方式中，numpy.count_nonzero 似乎是最快的:

相比之下，您的原始实现在该图上的最后一个输入大约需要 0.2 秒，因此比最差的 numpy 实现长约 20 倍(比最快的实现长约 100 倍)。