Pytorch张量——随机替换满足条件的值

Question

我有一个 mask 维度的 Pytorch 张量，

torch.Size([8, 24, 24])

具有唯一值，

> torch.unique(mask, return_counts=True)
(tensor([0, 1, 2]), tensor([2093, 1054, 1461]))

我希望将2的个数随机替换为0，这样张量中的唯一值和个数就变成了，

> torch.unique(mask, return_counts=True)
(tensor([0, 1, 2]), tensor([2500, 1054, 1054]))

我试过使用 torch.where 但没有成功。如何实现？

Answer 1

一种可能的解决方案是通过 view 和 numpy.random.choice 进行扁平化：

from numpy.random import choice

idx = torch.where(mask.view(-1) == 2)[0]  # get all indicies of 2 in flat tensor

num_to_change = 2500 - 2093 # as follows from example abow

idx_to_change = choice(idx, size=num_to_change, replace=False)

mask.view(-1)[idx_to_change] = 0

Pytorch张量——随机替换满足条件的值

Pytorch tensor - randomly replace values that meet condition

python

replace

pytorch

tensor