TensorFlow中如何获取数组前N大元素的索引？多标签分类需求

当然有啦！针对你提到的两种场景，TensorFlow都有直接可用的函数或者简单的实现方案，我给你详细拆解一下：

场景1：获取数组前N个最大元素的索引 #

你需要的功能可以直接用tf.math.top_k()函数实现，它会返回数组中前k个最大的值以及对应的索引（默认按从大到小排序）。

举个例子，对应你给出的输入：

import tensorflow as tf

predicted_array = tf.constant([0.4, 0.6, 0.7, 0.2, 0.9])
# 获取前2个最大元素的索引
values, indices = tf.math.top_k(predicted_array, k=2)
print(indices.numpy())  # 输出: [4 2]

这里返回的索引是按数值从大到小排列的（0.9对应索引4，0.7对应索引2）。如果你需要索引按原数组中的位置从小到大排序，可以再加一步排序：

sorted_indices = tf.sort(indices)
print(sorted_indices.numpy())  # 输出: [2 4]

场景2：处理独热编码的真实标签数组 #

独热编码数组里我们需要的是值为1的元素的索引，这时候可以用tf.where()来定位非零元素，再提取索引：

对应你的示例：

one_hot_array = tf.constant([0, 0, 1, 0, 1])
# 找到所有值为1的索引
non_zero_indices = tf.where(tf.equal(one_hot_array, 1))
# 拉平成一维数组
flat_indices = tf.reshape(non_zero_indices, [-1])
# 取前N个（这里N=2，刚好是所有非零索引）
result = flat_indices[:2]
print(result.numpy())  # 输出: [2 4]

如果你的独热标签可能有超过N个1的情况，直接切片[:N]就能得到前N个索引（按原数组顺序排列）。

合并场景的通用函数 #

如果你想写一个通用函数，同时处理两种情况（比如既可以处理预测概率数组，也可以处理独热标签），可以这样实现：

def get_top_n_indices(arr, n):
    if tf.reduce_max(arr) == 1.0 and tf.reduce_min(arr) == 0.0:
        # 判定为独热编码数组
        indices = tf.reshape(tf.where(tf.equal(arr, 1)), [-1])
        return indices[:n]
    else:
        # 处理概率数组，返回排序后的前n个索引
        _, indices = tf.math.top_k(arr, k=n)
        return tf.sort(indices)

# 测试预测数组
pred_arr = tf.constant([0.4, 0.6, 0.7, 0.2, 0.9])
print(get_top_n_indices(pred_arr, 2).numpy())  # [2 4]

# 测试独热数组
one_hot_arr = tf.constant([0, 0, 1, 0, 1])
print(get_top_n_indices(one_hot_arr, 2).numpy())  # [2 4]

内容的提问来源于stack exchange，提问作者sdiabr