PyTorch中为2D张量每行生成动态Top K元素掩码（K为张量）

需求说明 #

给定任意2D张量X，以及一个与X行数相同的张量K（每行对应一个K值），生成布尔掩码，其中每行中Top K个最大元素的位置标记为True，其余为False。

输入示例 #

X = torch.tensor([[0.6607, 0.1165, 0.0278, 0.1950],
                  [0.0529, 0.4607, 0.2729, 0.2135],
                  [0.3267, 0.0902, 0.4578, 0.1253]])
K = torch.tensor([2,3,1])

期望输出 #

tensor([[ True,  False, False, True],
        [ False, True,  True,  True],
        [ False, False, True, False]])

解决方案 #

直接遍历每行，利用PyTorch的topk函数获取每行Top K元素的索引，再将掩码对应位置设为True即可：

import torch

# 定义输入张量
X = torch.tensor([[0.6607, 0.1165, 0.0278, 0.1950],
                  [0.0529, 0.4607, 0.2729, 0.2135],
                  [0.3267, 0.0902, 0.4578, 0.1253]])
K = torch.tensor([2,3,1])

# 创建与X形状相同的全False掩码
mask = torch.zeros_like(X, dtype=torch.bool)

# 逐行处理
for row_idx in range(X.shape[0]):
    # 获取当前行Top K元素的索引
    _, topk_indices = X[row_idx].topk(K[row_idx].item())
    # 将掩码对应位置设为True
    mask[row_idx, topk_indices] = True

print(mask)

代码说明 #

1. torch.zeros_like(X, dtype=torch.bool)：创建和X形状一致、初始值全为False的布尔张量作为掩码。
2. X[row_idx].topk(K[row_idx].item())：对第row_idx行的元素取Top K个最大值，返回的第二个元素就是这些最大值对应的索引。
3. mask[row_idx, topk_indices] = True：将掩码中对应Top K元素的位置标记为True。

如果处理大规模数据，也可以通过向量化操作优化，但上述循环实现逻辑清晰、易于调试，在行数不多的场景下足够高效。

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