如何仅将torch.topk应用于张量的非零元？

Question

我想应用函数torch.topk，但只应用于张量的非零元素(即，在计数过程中不计算零元素)。

目前我这样做：

torch.topk(tensor.view(-1), k)

但这也考虑了变量张量中的零元素，并返回其中最大的一个。我应该怎么做才能在非零元素中得到topk？

Answer 1

# get the top k values in a tensor excluding zeros
top_vals = torch.topk(tensor.view(-1), k)[0]
mask = top_vals != 0
values = top_vals[mask]
print(values)

# Get Indices of Top K Values
indices = torch.nonzero(mask)
print(indices)

学分：this question and its answers