Как реализовать в pytorch - .unique numpy With(!) return_index = True?

В numpy.unique есть опция return_index=True — которая возвращает позиции уникальных элементов (первое вхождение — если несколько).

К сожалению, в torch.unique такой опции нет!

Вопрос: Каковы быстрые и простые способы получить индексы уникальных элементов?

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

В более общем плане моя проблема заключается в следующем: у меня есть два вектора v1, v2, и я хочу получить позиции тех элементов в v2, которых нет в v1, а также для повторяющихся элементов мне нужна только одна позиция. Уникальный Numpy с return_index = True сразу дает решение. Как это сделать в факеле? Если мы знаем, что вектор v1 отсортирован, можно ли его использовать для ускорения процесса?

🤔 А знаете ли вы, что...
В Python есть множество библиотек и фреймворков для разработки веб-приложений.


3
59
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]

P.S. На моем компьютере функция обрабатывает вектор размером 10 миллионов примерно за (1,1 ± 0,1) с.