Pytorch中的torch.where函數

首先我們看一下Pytorch中torch.where函數是怎樣定義的：

@overload
def where(condition: Tensor) -> Union[Tuple[Tensor, ...], List[Tensor]]: ...

torch.where函數的功能如下：

torch.where(condition, x, y):
condition：判斷條件
x：若滿足條件，則取x中元素
y：若不滿足條件，則取y中元素

以具體實例看一下torch.where函數的效果：

import torch
 
# 條件
condition = torch.rand(3, 2)
print(condition)
# 滿足條件則取x中對應元素
x = torch.ones(3, 2)
print(x)
# 不滿足條件則取y中對應元素
y = torch.zeros(3, 2)
print(y)
# 條件判斷后的結果
result = torch.where(condition > 0.5, x, y)
print(result)

結果如下：

tensor([[0.3224, 0.5789],
        [0.8341, 0.1673],
        [0.1668, 0.4933]])
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]])
tensor([[0., 0.],
        [0., 0.],
        [0., 0.]])
tensor([[0., 1.],
        [1., 0.],
        [0., 0.]])

可以看到torch.where函數會對condition中的元素逐一進行判斷，根據判斷的結果選取x或y中的值，所以要求x和y應該與condition形狀相同。