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,这分别对应不同的维度进行操作,本例中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