PyTorch tensor的scatter_函数

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

把src里面的元素按照index和dim参数给出的条件，放置到目标tensor里面，在这里是self。下面为了讨论方便，目标tensor和self在交换使用的时候，请大家知道，在这里指的是同一个tensor.

注意：这里self, index, src三个张量的纬度必须是一致的(但每个纬度上的size不一定一致，请大家体会)。
只有src是个例外，可以是标量，即单个数字。
这个时候，就是把这单个数字，根据参数的条件， 放置到self的不同位置。

那么怎么放呢？根据PyTorch的文档，对于一个3-D的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

由上面的公式很容易推断出，对于一个2-D的tensor，放置方法如下：

self[index[i][j]][j] = src[i][j]  # if dim == 0
self[i][index[i][j]] = src[i][j]  # if dim == 1

对于一个1-D的tensor，放置方法如下：

self[index[i]] = src[i]  # if dim == 0       

是不是有点晕？我们来解释一下。

1. 当dim为0的时候

我们把src里面的元素放置到self里面的时候，假设是放置src的第[i][j][k]个元素，那么放置到self里面的位置(三个纬度的值)分别如下：

• index[i][j][k]
• j
• k

对于第一个纬度的位置，就是把i，替换为index[i][j][k]

那么这里有个问题，如果index的size比src的size要小的话，怎么办？ 那就是对于在index里面，找不到的值，就不再处理，self里面原来是什么还是什么。

为了更加方便说明，这里假设src是1-D的，即一个1维数组，那么dim只有一个值可以设置，即0(当然也可以说是有两个值，-1也是可以的，但是-1和0实际上指的都是第一个纬度)。那么这个时候self和index按照上面的规则，也必须都是一维的(参见上面的注意)。那么我们直接来看一段示例代码和输出来进行解释：

a = torch.arange(1, 6).long()
print(a)
i = torch.LongTensor([4,3,2])
t = torch.zeros(10).long()
t.scatter_(0, i, a)
print(t)

输出为：

tensor([1, 2, 3, 4, 5])
tensor([0, 0, 3, 2, 1, 0, 0, 0, 0, 0])

可以看到这里，源tensor(a)是一个一维，包含5个元素。目标tensor(t)，在这里是一个10个元素的tensor，为了大家看得方便，我们先把所有元素设置为0，然后再把源tensor里面的元素搬过来放进目标tensor里面的时候，就很容易看到，被index tensor里面的信息所影响到的元素是非0的，如果没受到影响的是0。

这里源tensor只有5个元素，那么都搬过来，目标tensor(t)里面的元素也还是有10-5=5个元素是不会受到影响的，即为0。

那么为什么上面看到目标tensor里面的非0元素的个数只有3个，而不是5个(等于源tensor的个数)? 回顾一下对于3-D的tensor，当dim=0的时候，元素设置的公式：

self[index[i][j][k]][j][k] = src[i][j][k]  # if dim == 0

显然，对于1-D的tensor，上面的公式简化为：

self[index[i]] = src[i]  # if dim == 0       

因为这里index只有三个元素[4,3,2]，那么意味着，再把源tensor(a)里面的5个元素放置到目标tensor(t)的过程中，只有i取值为0,1,2的，才能使用index里面的值，其余2个(在a里面的位置分别为4,5)，就不再般src里面的元素了。我们来逐个元素说明一下：

• 当 i == 0时，self[index[0]] = src[0]，即self[4] = src[0]，也就是把src里面的第1个元素设置到self的第4个元素，这里src[0] 即是 a[0]，是1，而self[4]，即t[4]被设置为了1.
• 当 i == 1时，self[index[1]] = src[1]，即self[3] = src[1]，也就是把src里面的第2个元素设置到self的第3个元素，这里src[1] 即是 a[1]，是2，而self[3]，即t[3]被设置为了2.
• 当 i == 2时，self[index[2]] = src[2]，即self[2] = src[2]，也就是把src里面的第3个元素设置到self的第2个元素，这里src[2] 即是 a[2]，是3，而self[2]，即t[2]被设置为了3.
• 当 i == 3 和4的是，index里面已经没有对应的数值了，这些元素就不处理了。

2. 当dim为1的时候

说明，src，目标tensor和index都至少是2-D的，如果设置dim = 1，将会导致PyTorch报错。错误信息如下(对于1-D的index)：

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

对于一个3-D的tensor，我们把src里面的元素放置到self里面的时候，假设是放置src的第[i][j][k]个元素，那么放置到self里面的位置(三个纬度的值)分别如下：

• i
• index[i][j][k]
• k

对于第一个纬度的位置，就是i，元素在src里面的位置是什么，在self里面也是相同的。

对于第二个纬度的位置，就是j，元素在self里面的位置变成了index[i][j][k]。

那么同样地，如果index的size比src的size要小的话，怎么办？那就是对于在index里面，找不到的值，就不在处理，即self里面是什么值还是什么值，不会变化。

对于一个2-D的tensor，我们把src里面的元素放置到self里面的时候，假设是放置src的第[i][j]个元素，那么放置到self里面的位置(个纬度的值)分别如下：

• i
• index[i][j]

为了更加方便理解，这里假设src是2-D的，即一个2维数组，且dim==1的情况下：

a = torch.arange(1, 11).long().reshape(2,5)
print(a)
i = torch.LongTensor([[4], [3]])
t = torch.zeros(10).long().reshape(2, 5)
t.scatter_(1, i, a)
print(t)

输出如下：

tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
tensor([[0, 0, 0, 0, 1],
        [0, 0, 0, 6, 0]])

在这里，index里面只有两个元素，那么也就是最终会有两个元素的值从src里面取出，设置到a里面去。在index里面仅有的两个元素是index[0][0]和index[1][0]，这两个对应的src的元素是a[0][0]和a[1][0]，对应的目的tensor(t)里面的t[0][index[0][0]]和t[1][index[1][0]]元素，即t[0][4]将会被设置为a[0][0]，t[1][3]将会被设置为a[1][0]，即：

• t[0][4] = 1
• t[1][3] = 6

其他目的tensor(t)里面的值都不会变。

3. 当dim为2的时候

大家可以按照上面说明的规则，自己进行推导，就不在这里赘述了。

总结：

scatter或者scatter_函数的作用就是把src里面的元素按照index和dim参数给出的条件，放置到目标tensor里面去。index有几个元素，就会有几个元素被从src里面放到目标tensor里面，其余目标tensor里面的元素不受影响。