pytorch scatter直观理解

简单记录一下pytorch scatter_的理解， 官方解释在 https://pytorch.org/docs/master/tensors.html#torch.Tensor.scatter_ 

官方的例子如下，下面说说使用层面的直观理解。

>>> src = torch.arange(1, 11).reshape((2, 5))
>>> src
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
>>> index = torch.tensor([[0, 1, 2, 0]])
>>> torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src)
tensor([[1, 0, 0, 4, 0],
        [0, 2, 0, 0, 0],
        [0, 0, 3, 0, 0]])

scatter的第一个参数是0，意思是index中的值指的是第几行。这么理解，index是[0,1,2,0]，对应的是说src的第一行[1,2,3,4,5]对应到0矩阵的每一个元素。

比如，index中的第一个0，是说src[0,0]，即1，对应到零矩阵的第0列第0行，

index中的第二个1,是说src[0,1]对应的是零矩阵的第1列第1行，

index中的最后一个0，说的是src[0,3]对应的是零矩阵的第3列第0行

再看一个例子，这个例子scatter的第一个参数是1，意思是index中坐标是指的第几列。

>>> index = torch.tensor([[0, 1, 2], [0, 1, 4]])
>>> torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src)
tensor([[1, 2, 3, 0, 0],
        [6, 7, 0, 0, 8],
        [0, 0, 0, 0, 0]])

Index中一共有两行，意思是一共替换原始的3行5列的矩阵中的前两行。第1行是[0,1,4]意思是替换原始零矩阵的第1行的第0列，第一行第1列和第一行第4列。我们看到确实零矩阵第1行的这三个数是被替换过得。

替换的是什么值呢？就是src中的第1行的0,1,2三列，即6,7,8.