torch.gather()通俗理解

torch.gather(input, dim, index, out=None) → Tensor
input (Tensor) – 源张量
dim (int) – 索引的轴
index (LongTensor) – 聚合元素的下标
out (Tensor, optional) – 目标张量

>>> t = torch.Tensor([[1,2],[3,4]])

1  2
3  4
>>> torch.gather(t,1,torch.LongTensor([[0,0],[1,0]])

1  1
4  3
[torch.FloatTensor of size 2x2]


下面,我将对上述的代码进行详细的解析:
首先,我们要看dim=1 or 0,1bi这分别对应不同的维度进行操作,本例中dim=1表示在横向,所以索引就是列号。index的第一行为[0,0],第一个0索引到原来t中第一列,所以对应第1行第1列元素即为1,第二个0也索引到第1列,即第1行第1个元素;index的第二行为[1,0],其中1索引到原来t中第二列即为4,0索引到第一列即为3,所以最终的结果为:

1  1
4  3

 

猜你喜欢

转载自blog.csdn.net/zl1107604962/article/details/103825678