Pytorch——在一个Tensor中随机保留一个True，其他全为False

需求

• 让一个[false,false,false,false,true,true]的tensor，变成随机只保留一个true的tensor，且可以并行处理：

要将一个布尔类型的张量 [False, False, False, False, True, True] 变成随机只保留一个 True 的张量，并且可以并行处理，可以使用以下代码：

import torch
import random

# 创建原始的布尔类型张量
tensor = torch.tensor([False, False, False, False, True, True])

# 获取所有 True 的索引
true_indices = torch.nonzero(tensor).squeeze()

# 随机选择一个索引
selected_index = random.choice(true_indices.tolist())

# 创建新的张量，只保留随机选择的 True
new_tensor = torch.zeros_like(tensor)
new_tensor[selected_index] = True

# 打印新的张量
print(new_tensor)

运行上述代码将输出类似如下的结果，其中只有一个 True 被保留，其他位置都是 False：

tensor([False, False, False, False, False,  True])

通过使用 torch.nonzero() 函数获取原始张量中所有 True 的索引，然后使用 random.choice() 函数从索引中随机选择一个索引，最后创建一个新的张量，只在随机选择的索引处设置为 True，其他位置都是 False，从而实现了随机只保留一个 True 的操作。这种方法可以并行处理多个布尔类型张量。