torch.where()详解

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

功能

其实非常简单，但是好像没有搜到个靠谱的解答，所以只能我自己写一个了。

d = torch.where(a>3, b, c)

简单的说： a是一个tensor，显然这个a>3是对a的每个元素进行判断，那么a的每个元素对应的位置就是 True 或者是 False

d 是一个shape和 b，c 相同的tensor，也就是 b，c 的shape 也必须相同。 d中的元素从b或者c中取，即：取b（True）或者c （False）中的元素。

注意：a不需要和b，c一样的shape

运行个简单的例子就明白了：

a 的shape是 (2, 2) b, c 的shape是 （3, 2, 2） 那么a>3会返回 (2, 2) 的True or False矩阵对吧，

a = torch.range(1, 4).reshape(2, 2)
b = torch.stack([a,a,a], dim=0)
c = b - 2
d = torch.where(a>3, b, c)
print('a:', a)
print('b:', b)
print('c:', c)
print('d:', d)

输出：可以看到，因为a中满足大于3的元素下标为(1, 1)，所以d的(:, 1, 1)位置取b的(:, 1, 1)的值：4，而别的下标对应的值取c对应的值。

a: tensor([[1., 2.],
        [3., 4.]])
b: tensor([[[1., 2.],
         [3., 4.]],

        [[1., 2.],
         [3., 4.]],

        [[1., 2.],
         [3., 4.]]])
c: tensor([[[-1.,  0.],
         [ 1.,  2.]],

        [[-1.,  0.],
         [ 1.,  2.]],

        [[-1.,  0.],
         [ 1.,  2.]]])
d: tensor([[[-1.,  0.],
         [ 1.,  4.]],

        [[-1.,  0.],
         [ 1.,  4.]],

        [[-1.,  0.],
         [ 1.,  4.]]])

完事儿

官方其实介绍很短：

Return a tensor of elements selected from either x or y, depending on condition.

The operation is defined as:

\text{out}_i = \begin{cases} \text{x}_i & \text{if } \text{condition}_i \\ \text{y}_i & \text{otherwise} \\ \end{cases}outi​={xi​yi​​if conditioni​otherwise​

NOTE

The tensors condition, x, y must be broadcastable.

• Parameters

  • condition (BoolTensor) – When True (nonzero), yield x, otherwise yield y
  • x (Tensor or Scalar) – value (if x is a scalar) or values selected at indices where condition is True
  • y (Tensor or Scalar) – value (if y is a scalar) or values selected at indices where condition is False

• Returns

  A tensor of shape equal to the broadcasted shape of condition, x, y

• Return type