pytorch torch.gather 和 torch.scatter_

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

依据dim指定的轴，取出input中index指定处的数据,输出与index相同shape的tensor
input和index必须有相同的维数，对于不是指定的dim维，index的维度==input的维度。对于在指定的dim维，index的维度<=input的维度

import torch
#当指定维度为0时，以列为单位，每列的值按照index中每列指定的值输出
t = torch.tensor([[1,2], [3,4]])
torch.gather(t, 0, torch.tensor([[0, 0,], [1, 0,]]))

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

import torch
#在第0维，index.shape[0]<=t.shape[0],其它维上index.shape[d]==input.shape[d]
t = torch.tensor([[1,2,5], [3,4,7]])     #2*3
index= torch.tensor([[0, 0,], [1, 0,]])  #2*2
#torch.gather(t, 0, index) 报错
index= torch.tensor([[0, 0,1]])         #1*3
torch.gather(t, 0, index)

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

#当指定维度为1时，以行为单位，每行的值，按照index中每行指定的值输出
import torch
t = torch.tensor([[1,2], [3,4]])
torch.gather(t, 1, torch.tensor([[0, 0,], [1, 0,]]))

tensor([[1, 1],
[4, 3]])

import torch
t = torch.tensor([[1,2,5], [3,4,7]])
torch.gather(t, 1, torch.tensor([[0, 0,], [1, 0,]]))

tensor([[1, 1],
[4, 3]])

Tensor.scatter_(dim, index, src, reduce=None) → Tensor

可以认为是gather的逆操作
将src按reduce操作依次安放到tensor中index指定位置上
reduce=None表示直接将值放到对应位置上
reduce='add’表示将值与原位置上的值对应相加得到新值
reduce='multiply’表示将值与原位置上的值对应相乘得到新值

import torch
src = torch.arange(1, 11).reshape((2, 5))
print('src:',src)
index = torch.tensor([[0, 1, 2, 0],[1, 2, 0, 1]])
print('index:',index)
torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src)

src: tensor([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10]])

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

tensor([[1, 0, 8, 4, 0],
[6, 2, 0, 9, 0],
[0, 7, 3, 0, 0]])

import torch
src = torch.arange(1, 11).reshape((2, 5))
print('src:',src)
index = torch.tensor([[0, 1, 2], [0, 1, 4]])
print('index:',index)
torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src)

src: tensor([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10]])
index: tensor([[0, 1, 2],
[0, 1, 4]])
tensor([[1, 2, 3, 0, 0],
[6, 7, 0, 0, 8],
[0, 0, 0, 0, 0]])