torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
参数:
- input 被索引的tensor
- dim 索引沿着的维度
- index 索引的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
可以看出以下几点:
- 输入的index的每一个值都会被作为输入dim的index来取值
- index的ndims和input的ndims相同
- 结果的shape与index的shape相同
写一个二维版理解下:
out[i][j] = input[index[i][j]][j] # if dim == 0
out[i][j] = input[i][index[i][j]] # if dim == 1
可以看出隐藏的限制:
- dim=0,index的维度1应小于等于input的维度1
- dim=1, index的维度0应小于等于input的维度0
>>>p=torch.randn(4,2)
>>>p
tensor([[ 0.7786, 1.4472],
[ 0.6529, -0.0105],
[ 0.8745, 0.0016],
[-0.8376, -0.4966]])
>>>p.gather(0,torch.tensor([[1,0]]))
tensor([[0.6529, 1.4472]])
>>>p.gather(0,torch.tensor([[1,0,1]]))
RuntimeError: Size does not match at dimension 1 expected index [1, 3] to be smaller than src [4, 2] apart from dimension 0
>>>p.gather(0,torch.tensor([[1,0],[0,1],[1,1],[0,0],[1,1]]))
tensor([[ 0.6529, 1.4472],
[ 0.7786, -0.0105],
[ 0.6529, -0.0105],
[ 0.7786, 1.4472],
[ 0.6529, -0.0105]])
>>>p.gather(1,torch.tensor([[1,0],[0,1]]))
tensor([[ 1.4472, 0.7786],
[ 0.6529, -0.0105]])
>>>p.gather(1,torch.tensor([[1,0,1],[0,1,1]]))
tensor([[ 1.4472, 0.7786, 1.4472],
[ 0.6529, -0.0105, -0.0105]])
>>>p.gather(1,torch.tensor([[1],[0],[1],[1],[1]]))
RuntimeError: Size does not match at dimension 0 expected index [5, 1] to be smaller than src [4, 2] apart from dimension 1
torch.gather
是取,torch.scatter
则是放
Tensor.scatter(dim, index, src, reduce=None) → Tensor
以index的值作为tensor的dim维度的索引,从src里取值
self[index[i][j][k]][j][k] = src[i][j][k] # if dim == 0
self[i][index[i][j][k]][k] = src[i][j][k] # if dim == 1
self[i][j][index[i][j][k]] = src[i][j][k] # if dim == 2
可以看出有以下限制:
- tensor,index,src的ndims相同
- index.size(d) <= src.size(d)
- 当d!=dim index.size(d) <= tensor.size(d)
- 如果src为数值,则默认维度与tensor一致
>>>x = torch.zeros(5,3)
>>>index = torch.tensor([0,1,2,0,1])
>>>x.scatter_(1,index.view(-1,1),1)
tensor([[1., 0., 0.],
[0., 1., 0.],
[0., 0., 1.],
[1., 0., 0.],
[0., 1., 0.]])
>>>index2 = torch.tensor([[0,1],[1,2],[2,0],[0,1],[1,2]])
>>>x.scatter_(1,index2,1)
tensor([[1., 1., 0.],
[0., 1., 1.],
[1., 0., 1.],
[1., 1., 0.],
[0., 1., 1.]])
>>>index3 = torch.randint(1,3,(5,4))
>>>x.scatter_(1,index3,1)
Traceback (most recent call last):
File "<input>", line 1, in <module>
RuntimeError: Expected index [5, 4] to be smaller than self [5, 3] apart from dimension 1 and to be smaller size than src [5, 3]
>>>src = torch.ones(5,4)
>>>x.scatter_(1,index3,src)
tensor([[1., 1., 1.],
[0., 1., 1.],
[1., 1., 1.],
[1., 1., 1.],
[0., 1., 1.]])