Question

我有一个np.argmax（model.predict（X），axis = 2）格式的预测，该预测返回一个元素。如何使用numpy预测前k个元素

Answer 1

@desertnaut提供的链接涵盖了1D情况。但是，将the good answer泛化为“沿轴ND”并不是完全简单的。

在此示例中，我们找到了沿轴1的顶部2：

>>> a = np.random.randint(0, 9, (3, 5, 6))
>>> b = a.argpartition(-2, axis=1)[:, -2:]
>>> i, j, k = a.shape
>>> i, j, k = np.ogrid[:i, :j, :k]
>>> b = b[i, a[i, b, k].argsort(axis=1), k]
>>> a
array([[[8, 4, 1, 2, 4, 8],
        [0, 1, 3, 4, 2, 7],
        [4, 2, 7, 8, 1, 4],
        [1, 6, 2, 0, 3, 7],
        [1, 0, 0, 2, 8, 1]],

       [[1, 6, 3, 3, 0, 6],
        [7, 2, 0, 3, 8, 5],
        [5, 0, 1, 1, 7, 4],
        [2, 2, 4, 2, 6, 2],
        [5, 5, 7, 6, 8, 1]],

       [[4, 4, 4, 6, 2, 5],
        [2, 7, 8, 2, 6, 0],
        [5, 6, 7, 5, 1, 6],
        [6, 5, 3, 2, 2, 3],
        [5, 1, 8, 1, 6, 8]]])
>>> a[i, b, k]
array([[[4, 4, 3, 4, 4, 7],
        [8, 6, 7, 8, 8, 8]],

       [[5, 5, 4, 3, 8, 5],
        [7, 6, 7, 6, 8, 6]],

       [[5, 6, 8, 5, 6, 6],
        [6, 7, 8, 6, 6, 8]]])

一般功能看起来像

def argtopk(A, k, axis=0):
    tk = A.argpartition(-k, axis=axis)[(*axis*(slice(None),), slice(-k, None))]
    I = np.ogrid[(*map(slice, A.shape),)]
    I[axis] = tk
    I[axis] = A[I].argsort(axis=axis)
    return tk[I]

如何沿一个轴获取多维数组的k个顶级元素，而不是argmax给出的一个元素

1 个答案: