Pytorch系列（1）：torch.gather()

torch.gather

作用：收集输入的特定维度指定位置的数值

参数：

input(tensor):   待操作数。不妨设其维度为（x₁, x₂, …, x_n）
dim(int):   待操作的维度。
index(LongTensor):   如何对input进行操作。其维度有限定，例如当dim=i时，index的维度为（x₁, x₂, …y, …,x_n），既是将input的第i维的大小更改为y，且要满足y>=1（除了第i维之外的其他维度，大小要和input保持一致）。
out:   注意输出和index的维度是一致的

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

解释及举例

torch.gather()函数从公式上来说并不容易理解，我们以一个场景为例。

在序列标注问题上，我们给每一个单词都标上一个标签。不