如何在张量经过step function后，筛选原张量中top-k的阈值达标元素？

给定张量：

import torch
my_tensor = torch.tensor([1.0, -0.5, -0.2, 0.6, 0.88])

通过step函数（逻辑：超过阈值设为1，低于阈值设为0）得到输出：

values = torch.tensor([0.0])
step_func_out = torch.heaviside(my_tensor, values)
# 输出: torch.tensor([1.0, 0.0, 0.0, 1.0, 1.0])

需要基于原张量的值，从符合阈值要求（即step_func_out为1）的元素中选取top-k个，其余符合条件的元素设为0。例如：

• 当k=2时，输出为torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0])（选原张量中符合条件的前2大元素：1.0和0.88）
• 当k=4时，输出为torch.tensor([1.0, 0.0, 0.0, 1.0, 1.0])（符合条件的只有3个元素，全部保留）

可以通过以下步骤实现需求：

1. 筛选出符合阈值条件的元素及其原索引
2. 在符合条件的元素中找出top-k的位置（映射回原张量索引）
3. 初始化全0张量，将top-k索引位置设为1

代码实现：

import torch

def select_top_k_qualified(my_tensor, threshold=0.0, k=2):
    # 1. 获取符合阈值条件的元素掩码与原索引
    mask = my_tensor > threshold
    qualified_indices = torch.where(mask)[0]
    qualified_vals = my_tensor[mask]
    
    # 2. 确定实际要保留的元素数量，避免k超过符合条件的元素总数
    keep_num = min(k, len(qualified_vals))
    # 获取符合条件元素中top-k的位置，再映射回原张量索引
    _, top_k_pos = torch.topk(qualified_vals, keep_num)
    top_k_indices = qualified_indices[top_k_pos]
    
    # 3. 生成最终输出张量
    final_output = torch.zeros_like(my_tensor)
    final_output[top_k_indices] = 1.0
    return final_output

# 测试示例
my_tensor = torch.tensor([1.0, -0.5, -0.2, 0.6, 0.88])
print(select_top_k_qualified(my_tensor, k=2))
# 输出: tensor([1., 0., 0., 0., 1.])
print(select_top_k_qualified(my_tensor, k=4))
# 输出: tensor([1., 0., 0., 1., 1.])

• 掩码筛选：mask = my_tensor > threshold生成布尔张量标记符合条件的元素，torch.where(mask)[0]获取这些元素在原张量中的索引。
• top-k处理：用torch.topk从符合条件的元素中选出前keep_num个（keep_num取k和符合条件元素数的最小值，避免索引越界），再将这些位置映射回原张量的索引。
• 生成输出：创建与原张量形状一致的全0张量，将top-k索引对应的位置设为1，得到最终结果。

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