python - 随机获取 PyTorch 张量中最大值之一的索引

Question

我需要在一维张量上执行类似于内置 torch.argmax() 函数的操作，但我希望能够随机选择一个最大值而不是选择第一个最大值的索引最大值之一的索引。例如:

my_tensor = torch.tensor([0.1, 0.2, 0.2, 0.1, 0.1, 0.2, 0.1])
index_1 = random_max_val_index_fn(my_tensor)
index_2 = random_max_val_index_fn(my_tensor)

print(f"{index_1}, {index_2}")
> 5, 1

Answer 1

可以先获取所有最大值的索引，然后从中随机选择:

def rand_argmax(tens):
    max_inds, = torch.where(tens == tens.max())
    return np.random.choice(max_inds)

样本运行:

>>> my_tensor = torch.tensor([0.1, 0.2, 0.2, 0.1, 0.1, 0.2, 0.1])
>>> rand_argmax(my_tensor)
2
>>> rand_argmax(my_tensor)
5
>>> rand_argmax(my_tensor)
2
>>> rand_argmax(my_tensor)
1