首页 > temp > 简明python教程 >
-
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):
1
2
3
4
5
6
|
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
7
8
9
10
11
|
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 ] |