pytorch 深入理解 tensor.scatter_ ()用法

在 pytorch 库下理解 torch.tensor.scatter()的用法。作者在网上搜索了很多方法，最后还是觉得自己写一篇更为详细的比较好，转载请注明。
首先，scatter() 和 scatter_() 的作用是一样的，但是 scatter() 不会直接修改原来的 Tensor，而 scatter_() 会修改原先的 Tensor。

1 API格式

torch.Tensor.scatter_(dim, index, src) → Tensor 

	
	

		1
	

字面意思：对一个 torch.Tensor 进行操作，dim，index，src三个为输入的参数。

• dim 就是在哪个维度进行操作，注意，dim 的不同，在其他条件相同的条件下得到的output 也不同。
• index 是输入的索引。
• src 就是输入的向量，也就是 input。

最后，函数返回一个 Tensor。

2 具体示例

import torch as th
# import torch 包

a = th.rand(2,5)	
# 初始化向量 a，size 为 (2, 5)，二维向量，2行5列，每个元素是 0 到 1 的均匀分布采样
# 把 a 作为 src，也就是 input 
# a 的初始化数值如下: 
src tensor:
tensor([[0.6789, 0.7350, 0.6104, 0.7777, 0.9613],
        [0.1432, 0.8788, 0.3269, 0.0063, 0.6070]])

        
# 初始化 b 为size 为 (3, 5) 的向量，二维向量，3行5列，每个元素被初始化为 0
b = th.zeros(3, 5).scatter_(
	dim = 0,
	index = th.LongTensor([[0, 1, 2, 0, 0],[2, 0, 0, 1, 2]]),
	src = a 
)
# dim = 0, out:
tensor([[0.6789, 0.8788, 0.3269, 0.7777, 0.9613],
        [0.0000, 0.7350, 0.0000, 0.0063, 0.0000],
        [0.1432, 0.0000, 0.6104, 0.0000, 0.6070]])
        
# 初始化 c 为size 为 (3, 5) 的向量，二维向量，3行5列，每个元素被初始化为 0
c = th.zeros(3, 5).scatter_(
	dim = 1,
	th.LongTensor([[0, 1, 2, 0, 0],[2, 0, 0, 1, 2]]),
	src = a	
)
# dim = 1, out:
tensor([[0.9613, 0.7350, 0.6104, 0.0000, 0.0000],
        [0.3269, 0.0063, 0.6070, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]) 

下面来解释一下，b，c 内的元素分别是怎么得到的。

2.1 dim = 0 下的结果分析

先说 b，也就是 dim =0 下得到的结果。我们来看下官方给的说明文字：

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 

因为这时 dim = 0，且只有 2 个维度，所以我们只用看第一行就行。

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

仅用这一个公式就确定了 b 中所有元素的取值，与 a 的映射关系。这里等号左边的 self 可看做 output，也就是 b；src 是我们的输入向量，也就是 a。这里的 i，j 分别是输入向量 src 的 size 的取值。比如，本例中 a 的 size 为 (2，5)，也就是说，对于 a 中的元素，i 的取值为 0，1；j 的取值为 0，1，2，3，4。a 中的元素的索引也就是(0，0)，(0，1)，… (0，4)；(1，0)，(1，1)，…(1，4) 完毕，一共 2*5 = 10 个元素。
了解了这些以后，通过举例来说明 b 中的元素都是如何确定的。

index = th.LongTensor([[0, 1, 2, 0, 0],[2, 0, 0, 1, 2]])，
我们列举一些元素来说明其映射关系

当 i = 0，j = 0 时，
我们用类似上述确定 a 索引的方式确定了 index[i][j] = 0，
这里的 0 就是 [0，1，2，0，0] 中最左边的 0，
则 b = out[index[i][j]][j] = out[0][0] = src[0][0] = 0.6789

当 i = 0，j = 1 时，index[0][1] = 1，
这里的 1 就是 [0，1，2，0，0] 中的 1，
同理，b = out[index[i][j]][j] = out[1][1] = src[0][1] = 0.7350

当 i = 0，j = 2 时，index[0][2] = 2，
这里的 2 就是 [0，1，2，0，0] 中的 2，
同理，b = out[index[i][j]][j] = out[2][2] = src[0][2] = 0.6104
注意，这里的out[2][2] 不是第 2 行，第 2 列的元素，是第 3 行，第 3 列的元素

当 i = 1，j = 1 时，index[1][1] = 0，
这里的 0 就是 [2，0，0，1，2] 中最**左**边的 0，
同理，b = out[index[i][j]][j] = out[0][1] = src[1][1] = 0.8788

当 i = 1，j = 3 时，index[1][3] = 0，
这里的 0 就是 [2，0，0，1，2] 中最**右**边的 0，
同理，b = out[index[i][j]][j] = out[0][1] = src[1][3] = 0.0063

当 i = 1，j = 4 时，index[1][4] = 2，
这里的 2 就是 [2，0，0，1，2] 中最**左**边的 0，
同理，b = out[index[i][j]][j] = out[0][1] = src[1][4] = 0.6070

由此得到了 b 中有映射关系的元素，剩余的元素，由于 b 被初始化为全 0 向量，所以剩余的元素均为 0 。 

dim = 1的时候，同理。只是换了一种映射机制，如法炮制。