pytorch: where、gather函数

**

一、where函数

**
torch.where(condition,x,y)
out = x,if condition is 1
= y ,if condition is 0

In [29]: cond = torch.rand(2,2)

In [30]: cond
Out[30]:
tensor([[0.1326, 0.4126],
        [0.7093, 0.5339]])

In [31]: a = torch.zeros(2,2)

In [32]: a
Out[32]:
tensor([[0., 0.],
        [0., 0.]])

In [33]: b = torch.ones(2,2)

In [34]: b
Out[34]:
tensor([[1., 1.],
        [1., 1.]])

In [35]: torch.where(cond>0.5,a,b)
Out[35]:
tensor([[1., 1.],

**

二、gather函数

**
torch.gather(input,dim,index,out=None)

In [1]: import torch

In [2]: prob = torch.randn(4,10)

In [3]: prob
Out[3]:
tensor([[ 0.8383,  0.2332,  1.1231,  0.4929,  2.1630,  0.9328, -0.9775,  2.9904,
          1.7534,  0.2515],
        [-1.1460,  0.3640, -0.6829, -2.0924,  0.2590,  1.8114,  0.8341, -0.6201,
         -0.8322, -0.1316],
        [ 2.3255, -0.4369,  0.6470,  1.0118, -0.4143,  0.5650, -0.2035, -0.1714,
         -1.1866, -0.0068],
        [ 0.9149, -0.3263, -0.3857,  2.1448, -0.2073,  1.5119, -1.1741, -1.6586,
          0.1826,  1.0848]])

In [4]: idx = prob.topk(dim=1,k=3)

In [5]: idx = idx[1]

In [6]: idx
Out[6]:
tensor([[7, 4, 8],
        [5, 6, 1],
        [0, 3, 2],
        [3, 5, 9]])

In [7]: label = torch.arange(10)+100

In [8]: label
Out[8]: tensor([100, 101, 102, 103, 104, 105, 106, 107, 108, 109])

In [9]: torch.gather(label.expand(4,10),dim=1,index=idx.long())
Out[9]:
tensor([[107, 104, 108],
        [105, 106, 101],
        [100, 103, 102],
        [103, 105, 109]])