如何仅对Tensor的非零元素应用torch.topk函数？

可以通过**先过滤出所有非零元素，再对过滤后的张量应用torch.topk**来实现需求，具体步骤如下：

1. 提取张量中的非零元素
   使用布尔索引直接筛选出非零元素，得到仅包含非零值的一维张量：

   non_zero_elements = tensor[tensor != 0]
   

   或者使用torch.nonzero配合索引（效果一致）：

   non_zero_elements = tensor[tensor.nonzero(as_tuple=True)]
   

2. 对非零元素执行topk操作
   直接对过滤后的张量调用torch.topk即可，同时建议处理非零元素数量小于k的边界情况：

   # 实际取的k值为非零元素数量和输入k的较小值
   actual_k = min(k, len(non_zero_elements))
   if actual_k > 0:
       top_k_vals, top_k_indices = torch.topk(non_zero_elements, actual_k)
   else:
       # 无任何非零元素时的处理，比如返回空张量或默认值
       top_k_vals = torch.tensor([])
       top_k_indices = torch.tensor([])
   

完整示例代码 #

import torch

tensor = torch.tensor([[0, 5, 0], [3, 0, 8], [2, 7, 0]])
k = 3

# 提取非零元素
non_zero_elements = tensor[tensor != 0]
# 处理边界情况
actual_k = min(k, non_zero_elements.numel())
if actual_k > 0:
    top_k_vals, top_k_indices = torch.topk(non_zero_elements, actual_k)
else:
    top_k_vals = torch.tensor([])
    top_k_indices = torch.tensor([])

print("非零元素：", non_zero_elements)
print(f"前{actual_k}个最大值：", top_k_vals)
print("对应索引（在非零元素张量中的位置）：", top_k_indices)

如果需要获取这些最大值在原张量中的位置，可以在提取非零元素时同时记录原索引：

# 获取非零元素的原索引
non_zero_indices = torch.nonzero(tensor, as_tuple=True)
non_zero_elements = tensor[non_zero_indices]

actual_k = min(k, non_zero_elements.numel())
if actual_k > 0:
    top_k_vals, top_k_pos_in_nonzero = torch.topk(non_zero_elements, actual_k)
    # 映射回原张量的索引
    top_k_indices = tuple(idx[top_k_pos_in_nonzero] for idx in non_zero_indices)
else:
    top_k_vals = torch.tensor([])
    top_k_indices = (torch.tensor([]), torch.tensor([]))

print("前k个最大值在原张量中的索引：", top_k_indices)

内容的提问来源于stack exchange，提问作者K.N