gather函数

gather(input, dim, index)：根据 index，在 dim 维度上选取数据，输出的 size 与 index 一致

# input (Tensor) – 源张量

# dim (int) – 索引的轴

# index (LongTensor) – 聚合元素的下标(index需要是torch.longTensor类型)

# out (Tensor, optional) – 目标张量

for 3D tensor:

out[i][j][k] = tensor[index[i][j][k]][j][k] # dim=0

out[i][j][k] = tensor[i][index[i][j][k]][k] # dim=1

out[i][j][k] = tensor[i][j][index[i][j][k]] # dim=2

for 2D tensor:

out[i][j] = input[index[i][j]][j] # dim = 0

out[i][j] = input[i][index[i][j]] # dim = 1

import torch as t  # 导入torch模块
c = t.arange(0, 60).view(3, 4, 5)  # 定义tensor
print(c)
index = torch.LongTensor([[[0,1,2,0,2],
　　　　　　　　　　　　　　　　[0,0,0,0,0],
　　　　　　　　　　　　　　　　[1,1,1,1,1]],
　　　　　　　　　　　　　　　 [[1,2,2,2,2],
　　　　　　　　　　　　　　 　 [0,0,0,0,0],
　　　　　　　　　　　　　　　  [2,2,2,2,2]]])

b = t.gather(c, 0, index)
print(b)

输出：

tensor([[[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19]],

[[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29],
[30, 31, 32, 33, 34],
[35, 36, 37, 38, 39]],

[[40, 41, 42, 43, 44],
[45, 46, 47, 48, 49],
[50, 51, 52, 53, 54],
[55, 56, 57, 58, 59]]])

报错：

Traceback (most recent call last):
File "E:/Release02/my_torch.py", line 14, in <module>
b = t.gather(c, 0, index)
RuntimeError: Size does not match at dimension 1 get 4 vs 3

（第1维尺寸不匹配）

将index调整为：

index = t.LongTensor([[[0, 1, 2, 0, 2], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]],
                      [[1, 2, 2, 2, 2], [0, 0, 0, 0, 0], [2, 2, 2, 2, 2], [1, 1, 1, 1, 1]],
                      [[1, 2, 2, 2, 2], [0, 0, 0, 0, 0], [2, 2, 2, 2, 2], [1, 1, 1, 1, 1]]])
则上文输出为：

tensor([[[ 0, 21, 42, 3, 44],
[ 5, 6, 7, 8, 9],
[30, 31, 32, 33, 34],
[35, 36, 37, 38, 39]],

[[20, 41, 42, 43, 44],
[ 5, 6, 7, 8, 9],
[50, 51, 52, 53, 54],
[35, 36, 37, 38, 39]],

[[20, 41, 42, 43, 44],
[ 5, 6, 7, 8, 9],
[50, 51, 52, 53, 54],
[35, 36, 37, 38, 39]]])

对于2D tensor 则无“index与tensor 的size一致”之要求，

这个要求在官方文档和其他博文、日志中均无提到

（可能是个坑吧丨可能是个坑吧丨可能是个坑吧）

eg:

代码（此部分来自https://www.yzlfxy.com/jiaocheng/python/337618.html）：

b = torch.Tensor([[1,2,3],[4,5,6]])
print b
index_1 = torch.LongTensor([[0,1],[2,0]])
index_2 = torch.LongTensor([[0,1,1],[0,0,0]])
print torch.gather(b, dim=1, index=index_1)
print torch.gather(b, dim=0, index=index_2)

输出：

 1 2 3
 4 5 6
[torch.FloatTensor of size 2x3]

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

 1 5 6
 1 2 3
[torch.FloatTensor of size 2x3]




官方文档：

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

 Gathers values along an axis specified by dim.

 For a 3-D tensor the output is specified by:

 out[i][j][k] = input[index[i][j][k]][j][k] # dim=0
 out[i][j][k] = input[i][index[i][j][k]][k] # dim=1
 out[i][j][k] = input[i][j][index[i][j][k]] # dim=2

 Parameters: 

  input (Tensor)-The source tensor
  dim (int)-The axis along which to index
  index (LongTensor)-The indices of elements to gather
  out (Tensor, optional)-Destination tensor

 Example:

 >>> t = torch.Tensor([[1,2],[3,4]])
 >>> torch.gather(t, 1, torch.LongTensor([[0,0],[1,0]]))
  1 1
  4 3
 [torch.FloatTensor of size 2x2]


以上，学习中遇到的问题，记录方便回顾，亦示他人以之勉坑

猜你喜欢