PyTorch如何获取二维张量行维度大于阈值的TopK值与索引

核心逻辑依托torch.topk的排序特性做优化，不需要逐行遍历全量元素，性能和原生topk接口基本一致：

• 首先按行取前10大的元素及对应列索引，topk默认按降序返回结果，每行第一个元素就是该行全局最大值
• 逐行过滤前10个元素中大于设定阈值的项：
  • 若存在符合阈值要求的元素，直接保留这些过滤后的值和索引即可，不需要凑满10个
  • 若前10个元素里没有任何一个大于阈值，说明整行所有元素都小于阈值（因为前10个已经是该行最大的10个值），直接返回该行最大值和对应索引即可

import torch

# 初始化测试张量
tensor_3x100 = torch.rand(3, 100)
# 配置参数：阈值、取前k个
threshold = 0.95
k = 10

# 按行取前k大的元素值和对应索引
topk_values, topk_indices = torch.topk(tensor_3x100, k=k, dim=1)

row_results = []
for val_row, idx_row in zip(topk_values, topk_indices):
    # 生成阈值过滤掩码
    valid_mask = val_row > threshold
    filtered_vals = val_row[valid_mask]
    filtered_idxs = idx_row[valid_mask]
    
    if len(filtered_vals) != 0:
        # 存在符合阈值的元素，直接保留过滤结果
        row_results.append( (filtered_vals, filtered_idxs) )
    else:
        # 无符合阈值的元素，仅保留该行最大值
        row_results.append( (val_row[:1], idx_row[:1]) )

# 结果验证
for row_num, (vals, indices) in enumerate(row_results):
    print(f"行{row_num} 符合条件的值：{vals}")
    print(f"行{row_num} 对应列索引：{indices}\n")

• 该实现逻辑严谨：如果topk返回的前10个最大值里都没有超过阈值的元素，该行剩余的90个元素必然小于等于第10个值，不可能满足阈值要求，不需要额外遍历全量数据
• 返回结果以列表形式存储每行的（值、索引）元组，适配每行返回元素数量不一致的场景，如果需要张量化输出，可以根据业务需求对短行做填充处理
• 注意：你原有代码里直接对topk返回对象调用gt()是错误写法，topk接口返回的是包含values、indices两个字段的专有结构，需要先取values属性再做阈值判断

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