pytorch中scatter()、scatter_()详解

scatter()、scatter_()

scatter() 和 scatter_() 的作用一样。

不同之处在于 scatter() 不会直接修改原来的 Tensor，而 scatter_() 会在原来的基础上对Tensor进行修改。

torch.scatter()官方文档

scatter(dim, index, src)将src中数据根据index中的索引按照dim的方向进行填充。

dim=0

# 填充规则(dim=0)
self[ index[i][j] ] [j] = src[i][j]

dim=1

# 填充规则(dim=1)
self [i] [ index[i][j] ] = src[i][j]

注：index必须是torch.LongTensor类型

index行数必须和src行数相同，且下标不能超过output的列数。
index列数不受限制。

src列数必须要小于等于进行填充的tensor列数（或者只是一个标量也可以）。

scatter() 用于 one-hot 编码

这就是一个典型的用标量来修改张量的一个例子。

class_num = 10
batch_size = 4
label = torch.LongTensor(batch_size, 1).random_() % class_num
#tensor([[6],
#        [0],
#        [3],
#        [2]])
torch.zeros(batch_size, class_num).scatter_(1, label, 1)
#tensor([[0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
#        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
#        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
#        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]])