Pytorch 的 gather的直观理解(做个记录)

torch.gather(input, dim, index, out=None),注:pytorch要求,除了dim指定的维度外,input和index在其他维度上必须大小一致.

直观的通俗解释:

如果都是2维的,dim = 0, 表示从input中的每一列中,选择第index行的值,生成一个新tensor

                               dim = 1, 表示从input中的每一行中,选择第index列的值,生成一个新tensor

例如:

dim=1表示,从input中的每一行中,选择index=[[1],[0],[0]],即选择每一行的第1,0,0个元素,构成新tensor

input                                                                  
tensor([[0.2973, 0.6688],
        [0.6045, 0.5933],
        [0.6964, 0.2374]])

index                                                                  
tensor([[1],
        [0],
        [0]])

torch.gather(input,dim=1,index=index)                                   
tensor([[0.6688],
        [0.6045],
        [0.6964]])

dim=1表示,从input中的每一行中,选择index=[[1,0],[0,0],[0,1]],即选择第一行的第1,0个,第二行选择第0,0个,第三行选择第0,1个元素,构成新tensor

index                                                                  
tensor([[1, 0],
        [0, 0],
        [0, 1]])
                            
torch.gather(input,dim=1,index=index)                                   
tensor([[0.6688, 0.2973],
        [0.6045, 0.6045],
        [0.6964, 0.2374]])

dim=0表示,从input中的每一列中,选择index=[[1,0]],即选择每一列的第1,0个元素,构成新tensor

input                                                                  
tensor([[0.2973, 0.6688],
        [0.6045, 0.5933],
        [0.6964, 0.2374]])

index                                                                  
tensor([[1,