如何从PyTorch张量中提取无重复的Top K值及对应索引

torch.Tensor.topk原生不支持返回无重复的前k个取值，可通过「稳定排序+连续去重」的逻辑实现需求，代码如下：

import torch

def unique_topk(input: torch.Tensor, k: int, dim: int = -1, largest: bool = True):
    # 稳定排序保证相同取值的原始先后顺序不变，默认降序取最大的topk
    sorted_vals, sorted_indices = torch.sort(input, dim=dim, descending=largest, stable=True)
    # 去除排序后连续重复的取值，保留每个取值第一次出现的位置
    unique_vals, _, mask = torch.unique_consecutive(
        sorted_vals, 
        dim=dim, 
        return_inverse=False, 
        return_counts=False, 
        return_mask=True
    )
    # 过滤得到去重后取值对应的原始索引
    unique_indices = sorted_indices[mask]
    # 处理去重后总数量小于k的边界情况
    take_k = min(k, unique_vals.size(dim))
    # 沿指定维度取前k个结果
    slice_ = [slice(None)] * input.ndim
    slice_[dim] = slice(take_k)
    return unique_vals[tuple(slice_)], unique_indices[tuple(slice_)]

测试示例 #

你给出的示例运行效果如下：

input = torch.tensor([0.2,0.2,0.1])
k = 2
dim = 0

vals, indices = unique_topk(input, k, dim)
print(vals)      # 输出：tensor([0.2000, 0.1000])
print(indices)   # 输出：tensor([0, 2])

实现说明 #

• 先通过torch.sort做稳定排序，排序后相同取值会连续排列，同时保留相同取值的原始先后顺序，确保我们能拿到每个取值第一次出现的索引
• 用torch.unique_consecutive去除连续重复的取值，相比全局去重的torch.unique性能更高，也更适配排序后的场景
• 支持任意维度的张量，指定对应dim参数即可；如果需要取最小的k个无重复值，将largest参数设为False即可
• 自动处理去重后总数量小于k的边界场景，此时会返回所有去重后的结果

内容的提问来源于stack exchange，提问作者namespace-Pt