PyTorch笔记之scatter()函数

scatter(dim, index, src) 的参数有 3 个

• dim：沿着哪个维度进行索引
• index：用来 scatter 的元素索引
• src：用来 scatter 的源元素，可以是一个标量或一个张量

直白点说就是在dim维度上以index将src中的数放到一个新的Tensor中。

首先说一下官网的计算规则：

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

从例子中可以看出，我们做的就是一个赋值操作，只不过不是直接self和src对应位置的赋值，而是根据index的值来在self的对应位置赋值。

另外需要注意，index中的值是不应大于该dim的通道数的。例如self是（3, 4, 5）的Tensor，我们选择的dim是2，则index的值必须小于5。

下面是两个例子：

1、one-hot编码（src是标量的情况下）：

ind = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
y = torch.zeros(3, 5).scatter_(0, ind, 1)

print(y)

'''
输出：
tensor([[1., 1., 1., 0., 0.],
        [1., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0.]])
'''

详细解释：

我们可以将该函数看作两个过程：

首先进行坐标的映射，根据公式 self[index[i][j][k]][j] = src[i][j]，所以其实这个坐标中j是不变的(绿色值)，我们只需要把index对应位置的值填入作为新的i即可。表中红色的值是坐标。

然后在y的对应位置填入1即可。（坐标这个多画了一行，请忽略）

 

2、 src是张量的情况下： 

x = torch.rand(2, 5)

ind = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
y = torch.zeros(3, 5).scatter_(0, ind, x)

print(x)
print(y)

'''
输出：
tensor([[0.1161, 0.8937, 0.7539, 0.9041, 0.5330],
        [0.7159, 0.1167, 0.0778, 0.9715, 0.7562]])
tensor([[0.1161, 0.1167, 0.0778, 0.9041, 0.5330],
        [0.0000, 0.8937, 0.0000, 0.9715, 0.0000],
        [0.7159, 0.0000, 0.7539, 0.0000, 0.7562]])
'''

详细解释：

与示例1类似，我们先得到对应的坐标映射，然后将src中(i, j)对应数据填到self的(index[i][j], j)中。