torch.gather用法笔记

gather的作用是在指定维度选择数据

import torch
#定义一个张量
x = torch.tensor([[1,2,3],[4,5,6]])
print(x)
#定义dim = 0 时的索引index0
index0 = torch.LongTensor([[0,1,1]])
print(index0)
y = torch.gather(x, 0, index0)
print(y)
#定义dim = 1 时的索引index1
index1 = torch.LongTensor([[0,1,1],[1,1,1]])
print(index1)
z = torch.gather(x, 1, index1)
print(z)

用法为torch.gather(input,dim,index)

input为定义的张量

tensor([[1, 2, 3],
        [4, 5, 6]])

dim为指定的维度，dim=0固定行，dim=1固定列，注意这里固定是index的维度，例如index0指定dim=0时，实际上index0是行维固定，列维待确定的索引

index0
(0,) (1,) (1,)

当执行gather后，index0的列维固定，此时index0为

index0
(0, 0) (1, 1) (1, 2)

在指定的张量x中按照index0索引选择数据，得到y

tensor([[1, 5, 6]])

index1指定dim=1时，实际上index1是列维固定，行维待确定的索引，此时此时index1为

index1
(,0),(,1),(,1)
(,1),(,1),(,1)

当执行gather后，index1的行维固定，此时index1为

index1
(0,0),(0,1),(0,1)
(1,1),(1,1),(1,1)

在指定的张量x中按照index1索引选择数据，得到z

tensor([[1, 2, 2],
        [5, 5, 5]])