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]
dim = 1时表示跨越列维度来对行维度进行映射。index_1中的[0, 1]代表选取b中的第一行列标为0, 1的元素。index_1中的[2,0]代表选取b中第二行列标为2,0的元素进行映射。
同理dim = 0时表示跨越行对列维度进行映射,index_2中第0列表示对b第0列的下标为0的行进行映射,index_2第1列表示对b第1列第1行和第0行进行映射。index_2第2列表示对b第2列的1行和0行进行映射。
import torch
import numpy as np
a = torch.LongTensor(np.arange(24)).view(2,3,4)
print(a)
'''
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]]])
'''
index = torch.LongTensor([[[0 ,1 ,2 ,0],
[0, 0, 0 ,0],
[1, 1, 1, 1]],
[[2, 2, 2, 2],
[1, 1, 1, 1],
[0, 0, 0, 0]]])
b = torch.gather(a, 1, index)
print(b)
"""
tensor([[[ 0, 5, 10, 3],
[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],
[[20, 21, 22, 23],
[16, 17, 18, 19],
[12, 13, 14, 15]]])
"""
我们设定dim=1,这时候映射的情况下,改变的只有第1个维度。对于例子中这个三维的张量a,对a[0]和a[1]两块,我们分别取出a[0]和a[1]的每一列,也就是[0,4,8]、[1,5,9]……[15,19,23]。每一列值之间第0个维度的索引和第2个维度的索引都是固定不变的,而改变的,只有dim=1的维度。这也是我们设定的。
同样的,取出index的这些列,也就是[0,0,1]、[1,0,1]……[2,1,0],这些值恰巧就是 映射后的张量 每一列值的对映射前每一列值的索引(表示不太清楚)。
具体的,我们拿出张量b的这些列(b是映射后的张量),也就是[0,0,4]、[5,1,5]……[23,19,15]
我们拿出a,index,b,的第一列。分别是[0,4,8],[0,0,1],[0,0,4],可以发现,index的列对a列的索引,刚好就是b的列。其余列也是一一对应。
3维部分引用CSDN博主「山坡上幼稚狗」
原文链接:https://blog.csdn.net/ccwlisha/article/details/101524389