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.]]])

完事儿