Pytorch 学习（5）：Pytorch中的 torch.gather/scatter_ 聚集/分散操作

torch.gather(input, dim, index, out=None) → Tensor 聚集操作
torch.Tensor scatter_(dim, index, src) → Tensor 分散操作

1.torch.gather(input, dim, index, out=None) → Tensor 聚集操作

沿给定轴 dim ,将输入索引张量 index 指定位置的值进行聚合，对一个 3 维张量,输出可以定义为:

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

如果 input 是 size 为 (x0,x1...,xi−1,xi,xi+1,...,xn−1) 且 dim = i 的 n 维张量,则 index 必须是具有 size 为 (x0,x1,...,xi−1,y,xi+1,...,xn−1) 的 n 维张量,其中 y >= 1 ,并且 out 将与 index 的 size 相同.

Parameters:	input (Tensor) – 源张量 dim (int) – 索引的轴 index (LongTensor) – 聚合元素的下标 out (Tensor, optional) – 目标张量

Example:

>>> t = torch.Tensor([[1,2],[3,4]])
>>> torch.gather(t, 1, torch.LongTensor([[0,0],[1,0]]))
 1  1
 4  3
[torch.FloatTensor of size 2x2]

聚集操作计算过程，这里维度等于1，从列的维度获取index中的位置，获取每一行元素。

out[0][0]=t[0][index[0][0]] =1    index[0][0]=0
out[0][1]=t[0][index[0][1]] =1    index[0][1]=0
out[1][0]=t[1][index[1][0]] =4    index[1][0]=1
out[1][1]=t[1][index[1][1]] =3    index[1][0]=0

2.torch.Tensor scatter_(dim, index, src) → Tensor分散操作

从张量src中按照index张量中指定的索引位置写入self张量的值。对于一个三维张量，self更新为：
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

参数：
input (Tensor): the source tensor
dim (int): the axis along which to index
index (LongTensor): the indices of elements to scatter
src (Tensor or float): the source element(s) to scatter

例子：

Example::
 
    >>> x = torch.rand(2, 5)
    >>> x
    tensor([[ 0.3992,  0.2908,  0.9044,  0.4850,  0.6004],
            [ 0.5735,  0.9006,  0.6797,  0.4152,  0.1732]])
    >>> torch.zeros(3, 5).scatter_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
    tensor([[ 0.3992,  0.9006,  0.6797,  0.4850,  0.6004],
            [ 0.0000,  0.2908,  0.0000,  0.4152,  0.0000],
            [ 0.5735,  0.0000,  0.9044,  0.0000,  0.1732]])
 
    >>> z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
    >>> z
    tensor([[ 0.0000,  0.0000,  1.2300,  0.0000],
            [ 0.0000,  0.0000,  0.0000,  1.2300]])

分散操作计算过程，这里维度等于0，对于一个行的索引i，从index中查找对应行索引值i的index的索引i,j，且index的列等于self[i][j]的j列，根据src[i][j]获取每一个元素。如查询无对应的索引位置，则仍保持self的值。

self[0][0]=self[index[0][0]][0] =src[0][0]=0.3992 
self[1][0]=self[index[NAN][0]][0] =src[NAN][0]=0.0000 
self[2][0]=self[index[1][0]][0] =src[1][0]=0.5735
......
self[2][2]=self[index[0][2]][2] =src[0][2]=0.9044

>>> z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
    >>> z
    tensor([[ 0.0000,  0.0000,  1.2300,  0.0000],
            [ 0.0000,  0.0000,  0.0000,  1.2300]])

计算过程：
self[0][0]=self[0][index[0][NAN]] =0.0000 
self[0][1]=self[0][index[0][NAN]] =0.0000 
self[0][2]=self[0][index[0]index[0][0]] =src[0][0]=1.23

Pytorch 学习（5）：Pytorch中的 torch.gather/scatter_ 聚集/分散操作

猜你喜欢