如何在pytorch中实现-numpy的.unique with(!) return_index = True？

Question

在 numpy.unique 中，有一个选项 return_index=True - 返回唯一元素的位置（如果有多个，则第一次出现）。

不幸的是，torch.unique 中没有这样的选项！

问题：有哪些快速且火炬式的方法来获取唯一元素的索引？

======================

更一般地说，我的问题如下：我有两个向量 v1、v2，我想获取 v2 中这些元素的位置，这些元素不在 v1 中，并且对于重复元素，我只需要一个位置。 Numpy 的 return_index = True 的唯一性立即给出了解决方案。 如何在火炬中做到这一点？如果我们知道向量 v1 是有序的，可以用它来加速这个过程吗？

Answer 1

您可以通过以下方法在 PyTorch 中实现此目的：

def get_unique_elements_first_idx(tensor):
    # sort tensor
    sorted_tensor, indices = torch.sort(tensor)
    # find position of jumps
    unique_mask = torch.cat((torch.tensor([True]), sorted_tensor[1:] != sorted_tensor[:-1]))
    return indices[unique_mask]

使用示例：

v1 = torch.tensor([2, 3, 3])
v2 = torch.tensor([1, 2, 6, 2, 3, 10, 4, 6, 4])

# Mask to find elements in v2 that are not in v1
mask = ~torch.isin(v2, v1)
v2_without_v1 = v2[mask]

# Get unique elements and their first indices
unique_indices = get_unique_elements_first_idx(v2_without_v1)

print(unique_indices) #[0, 3, 1, 2]
print(v2[mask][unique_indices ]) #[1, 4, 6, 10]