原文链接请参考:https://dreamhomes.top/posts/201906081516.html
参考官网:torch.gather
用法:torch.gather(input, dim, index, out=None, sparse_grad=False) → Tensor
收集输入的特定维度dim
指定位置index
的数值。
对于一个三维tensor,结果如下:
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
参数说明:
input (Tensor)
– 输入张量dim (int)
– 索引维度index (LongTensor)
– 收集元素的索引out (Tensor, optional)
– 目标张量sparse_grad (bool,optional)
– 输入为稀疏张量
直接看官网解释有点不明白,参考另一文章的实例说明:https://blog.csdn.net/cpluss/article/details/90260550
在序列标注问题上,我们给每一个单词都标上一个标签。不妨假设我们有4个句子,每个句子的长度不一定相同,标签如下:
input = [
[2, 3, 4, 5],
[1, 4, 3],
[4, 2, 2, 5, 7],
[1]
]
上例中有四个句子,长度分别为4,3,5,1,其中第一个句子的标签为2,3,4,5。我们知道,处理自然语言问题时,一般都需要进行padding,即将不同长度的句子padding到同一长度,以0为padding,那么上述经padding后变为:
input = [
[2, 3, 4, 5, 0, 0],
[1, 4, 3, 0, 0, 0],
[4, 2, 2, 5, 7, 0],
[1, 0, 0, 0, 0, 0]
]
那么问题来了,现在我们想获得每个句子中最后一个词语的标签,该怎么得到呢?既是,第一句话中的5,第二句话中的3,第三句话中7,第四句话中的1。
此时就需要用gather函数。
此时我们的input就是填充之后的tensor,dim=1, index就是各个句子的长度,即[[4],[3],[5],[1]]。之所以维度是4*1,是为了满足index维度和input维度之间的关系(讲解参数时有讲)。
代码如下所示:
In [26]: import torch
...: input = [
...: [2, 3, 4, 5, 0, 0],
...: [1, 4, 3, 0, 0, 0],
...: [4, 2, 2, 5, 7, 0],
...: [1, 0, 0, 0, 0, 0]
...: ]
...: input = torch.tensor(input)
...: #注意index的类型
...: length = torch.LongTensor([[4],[3],[5],[1]])
...: #index之所以减1,是因为序列维度是从0开始计算的
...: out = torch.gather(input, 1, length-1)
...: out
Out[26]:
tensor([[5],
[3],
[7],
[1]])