在张量列表中找到 n 个最小值

Question

我试图在 pytorch 的张量列表中找到 n 个最小值的索引。由于这些张量可能包含许多非唯一值，因此我不能简单地计算百分位数来获得索引。然而，非唯一值的顺序并不重要。

我想出了以下解决方案，但想知道是否有更优雅的方法：

import torch

n = 10

tensor_list = [torch.randn(10, 10), torch.zeros(20, 20), torch.ones(30, 10)]
all_sorted, all_sorted_idx = torch.sort(torch.cat([t.view(-1) for t in tensor_list]))

cum_num_elements = torch.cumsum(torch.tensor([t.numel() for t in tensor_list]), dim=0)
cum_num_elements = torch.cat([torch.tensor([0]), cum_num_elements])

split_indeces_lt = [all_sorted_idx[:n] < cum_num_elements[i + 1] for i, _ in enumerate(cum_num_elements[1:])]
split_indeces_ge = [all_sorted_idx[:n] >= cum_num_elements[i] for i, _ in enumerate(cum_num_elements[:-1])]
split_indeces = [all_sorted_idx[:n][torch.logical_and(lt, ge)] - c for lt, ge, c in zip(split_indeces_lt, split_indeces_ge, cum_num_elements[:-1])]

n_smallest = [t.view(-1)[idx] for t, idx in zip(tensor_list, split_indeces)]

编辑：理想情况下，解决方案会选择非唯一值的随机子集，而不是选择列表的第一个张量的条目。

Answer 1

Pytorch 确实提供了一种更优雅（我认为）的方式，使用 torch.unique_consecutive（参见 here）

我要处理张量，而不是张量列表，因为就像你自己做的那样，只有 cat 要做。之后解开索引也不难。

# We want to find the n=3 min values and positions in t 
n = 3
t = torch.tensor([1,2,3,2,0,1,4,3,2])
# To get a random occurrence, we create a random permutation
randomizer = torch.randperm(len(t))
# first, we sort t, and get the indices
sorted_t, idx_t = t[randomizer].sort()
# small util function to extract only the n smallest values and positions
head = lambda v,w : (v[:n], w[:n])
# use unique_consecutive to remove duplicates
uniques_t, counts_t = head(*torch.unique_consecutive(sorted_t, return_counts=True))
# counts_t.cumsum gives us the position of the unique values in sorted_t
uniq_idx_t = torch.cat([torch.tensor([0]), counts_t.cumsum(0)[:-1]], 0)
# And now, we have the positions of uniques_t values in t : 
final_idx_t = randomizer[idx_t[uniq_idx_t]]
print(uniques_t, final_idx_t)

>>> tensor([0,1,2]), tensor([4,0,1])
>>> tensor([0,1,2]), tensor([4,5,8])
>>> tensor([0,1,2]), tensor([4,0,8])

编辑：我认为添加的排列可以解决您的需要随机出现问题