torch.gather

torch.gather

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

参数：

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

 

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)
print(out)

tensor([[5],
        [3],
        [7],
        [1]])