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形状相同。

©️2020 CSDN 皮肤主题: 书香水墨 设计师:CSDN官方博客 返回首页