torch.where()理解

PyTorch中的torch.where()函数详解与应用

最新推荐文章于 2025-07-02 22:01:18 发布

原创最新推荐文章于 2025-07-02 22:01:18 发布 · 1.8k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

pytorch 专栏收录该内容

30 篇文章

订阅专栏

本文详细介绍了PyTorch中的torch.where()函数，该函数根据提供的条件从两个张量中选择元素。在第一种用例中，当条件满足时，函数返回张量x的对应元素，否则返回y的元素。在第二种用例中，torch.where()等同于torch.nonzero()，返回条件张量中非零元素的索引。通过示例展示了如何在二维和三维张量上使用该函数，用于数据筛选和索引操作。

部署运行你感兴趣的模型镜像

1. torch.where(condition, x, y) → Tensor

返回一个tensor，其中的元素根据condition从x或y中选取
$\text { out }_{i}= \begin{cases}\mathrm{x}_{i} & \text { if } \text { condition }_{i} \\ \mathrm{y}_{i} & \text { otherwise }\end{cases}$

case 1:

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
        [ 0.3898, -0.7197],
        [ 0.0478, -0.1657]])
>>> torch.where(x > 0, x, y)
tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])
>>> x = torch.randn(2, 2, dtype=torch.double)
>>> x
tensor([[ 1.0779,  0.0383],
        [-0.8785, -1.1089]], dtype=torch.float64)
>>> torch.where(x > 0, x, 0.)
tensor([[1.0779, 0.0383],
        [0.0000, 0.0000]], dtype=torch.float64)

2. torch.where(condition) → tuple of LongTensor

这种情况下torch.where完全等同于torch.nonzero(condition, as_tuple=True)，即找出condition不等于0的元素的索引，当condition是二维tensor时，返回一个n*2的tensor，condition是三维时，返回n*3的tensor.

case2:

torch.nonzero(torch.tensor([1, 1, 1, 0, 1]), as_tuple=True)
(tensor([0, 1, 2, 4]),)

>>> a = torch.randn(4,3)
>>> a
tensor([[-0.1272,  1.0898,  0.1808],
        [-0.6453, -1.5093,  2.1270],
        [ 0.3852,  1.0197, -0.2998],
        [-0.2550,  1.2353, -0.1208]])
>>> b = (a>0.1)
>>> b
tensor([[False,  True,  True],
        [False, False,  True],
        [ True,  True, False],
        [False,  True, False]])
>>> b.nonzero()
tensor([[0, 1],
        [0, 2],
        [1, 2],
        [2, 0],
        [2, 1],
        [3, 1]])

>>> a = torch.randn(2,3,4)
>>> a
tensor([[[ 0.0642, -2.1055, -1.0024,  0.6682],
         [-0.1130, -0.5724,  0.4137, -0.4402],
         [ 1.3719,  0.7411,  1.7062, -0.9893]],
 
        [[-0.1432,  0.4898, -0.4913,  1.1397],
         [-0.0883,  1.1581, -0.4784,  0.0195],
         [ 1.3767,  2.2682, -0.6881,  1.9919]]])
 
>>> b = (a>0.1)
>>> b
tensor([[[False, False, False,  True],
         [False, False,  True, False],
         [ True,  True,  True, False]],
 
        [[False,  True, False,  True],
         [False,  True, False, False],
         [ True,  True, False,  True]]])
>>> b.nonzero()
tensor([[0, 0, 3],
        [0, 1, 2],
        [0, 2, 0],
        [0, 2, 1],
        [0, 2, 2],
        [1, 0, 1],
        [1, 0, 3],
        [1, 1, 1],
        [1, 2, 0],
        [1, 2, 1],
        [1, 2, 3]])