【pytorch】scatter的使用

作用

scatter是“散开”的意思，顾名思义，是将一个Tensor按照index做分散。

形式

在pytorch中，scatter可以通过torch.scatter和torch.scatter_(修改自身数据)，或者Tensor自生就有的方法scatter

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

参数

1. input
   输入参数，如果是通过Tensor直接调用的，没有该参数（就是自身嘛），仅仅在torch.scatter/torch.scatter_中需要指定

2. dim
   维度，需要用于在哪一个维度

3. index
   为Tensor做scatter时需要的索引，注意，这个index必须是int64型。另外它的-1维必须与input的-1维一致。

4. src
   用于“分散”的源数据，它的shape必须与input的shape一致
   它将用于具体的计算，公式如下
   

5. value
   一个0-1.0之间的float数，它与src只要一个即可。如果指定了value，则按照value这个值去填充index指定的位置。

使用样例

一维数据

input = torch.zeros(3)
index = torch.LongTensor([0,2,1])
input = input.scatter(dim=0, index=index, src=torch.Tensor([10,20,30]))
print(input)

输出结果为，

tensor([10., 30., 20.])

怎么的出来的呢？
套用公式如下：
self[index[i]] = src[i]
则有

input[index[0]] = src[0]     =>   input[0] = 10
input[index[1]] = src[1]     =>   input[2] = 20
input[index[2]] = src[2]     =>   input[1] = 30

一维数据,重复index

input = torch.zeros(3)
index = torch.LongTensor([0,2,0])
input = input.scatter(dim=0, index=index, src=torch.Tensor([10,20,30]))
print(input)

结果为

tensor([30.,  0., 20.])

计算如下：

input[index[0]] = src[0]     =>   input[0] = 10
input[index[1]] = src[1]     =>   input[2] = 20
input[index[2]] = src[2]     =>   input[0] = 30   /*override here*/

一维数据，使用value

input = torch.zeros(3)
index = torch.LongTensor([0,2])
input = input.scatter(dim=0, index=index, value=0.999)
print(input)

结果如下

tensor([0.9990, 0.0000, 0.9990])

运算过程同上面例子，这里不再赘述

二维数据

x = torch.zeros(2,4)
index = torch.LongTensor([1,0,2,3]).reshape(1,4)
x = x.scatter(dim=1, index=index, src=torch.Tensor([10,20,30,40,50,60,70,80]).reshape(2,4))
print(x)

这个例子里，index是(1,4)，而input是(2,4)，这样仅仅针对第一"行"做scatter。
结果如下

tensor([[20., 10., 30., 40.],
        [ 0.,  0.,  0.,  0.]])

计算过程如下：

# row 0
input[0,index[0]] = src[0][0]     =>   input[0][1] = 10
input[0,index[1]] = src[0][1]     =>   input[0][0] = 20
input[0,index[2]] = src[0][1]     =>   input[0][2] = 30
input[0,index[3]] = src[0][1]     =>   input[0][3] = 40
# row 1
input[1,index[0]] = src[1][0]     =>   input[1][1] = 0
input[1,index[1]] = src[1][1]     =>   input[1][0] = 0
input[1,index[2]] = src[1][1]     =>   input[1][2] = 0
input[1,index[3]] = src[1][1]     =>   input[1][3] = 0