如何使用 np.cumsum 加快平均精度计算？

Question

我有，

scores = np.array([[0.9, 0.8, 0.6, 0.5, 0.4], [0.5, 0.4, 0.31, 0.21, 0.4 ]])
labels = np.array([[1, 0, 1, 1, 0], [0, 0, 0, 1, 1]])

我想在 K map 计算，我写了一个如下的算法，

        k=3
        mean_ap = 0
        n = len(scores)
        for i in range(n):
            cum = ap = 0.0
            idx = np.argsort(-scores[i])
            used_label = labels[i][idx][:k]
            m = sum(labels[i])
            for j, label in enumerate(used_label):
                cum += label
                ap += cum * label / (j + 1)
            mean_ap += ap / min(m, k)
        val = mean_ap / n

它基本上给出了如下计算公式： (1 + 0 + 2 / 3) / 3 + ( 0 + 0 + 1 / 3) / 2

有什么建议我可以使用 np.cumsum 来加速我的算法吗？ 我认为它已经过优化，我在这里看不到任何增强空间？

提前致谢。

Answer 1

希望这可以帮助你（我尽量避免 for 循环）：

k = 3
n = len(scores)
m = labels.sum(axis=1)

idx = np.argsort(-scores)
used_label = labels[:,idx][np.arange(0,n),np.arange(0,n),:k]

val = (np.cumsum(used_label, axis=1)*used_label / 
       np.arange(1,k+1) / 
       np.min([m,np.repeat(k,n)],axis=0).reshape(-1,1)).sum(axis=1).sum() / n