torch.scatter笔记

scatter() 和 scatter_() 的作用是一样的，只不过 scatter() 不会直接修改原来的 Tensor，而 scatter_() 会。

import torch
#定义一个张量
x = torch.tensor([[1,2,3],[4,5,6]])
print(x)
#定义一个3×3零张量
z = torch.zeros((3,3),dtype = int)
print(z)
#定义索引
index=torch.LongTensor([[0,1,1],[1,2,1]])
print(index)
z.scatter_(0,index,x)
print(z)

用法为torch.scatter_(dim,index,iuput)

input为定义的张量

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

dim为指定的维度，dim=0固定行，dim=1固定列，注意这里固定是index的维度，例如index0指定dim=0时，实际上index0是行维固定，列维待确定的索引

index0
(0,),(1,),(1,)
(1,),(2,),(1,)

当执行scatter_()后，index0的列维固定，此时index0为

index0
(0,0),(1,1),(1,2)
(1,0),(2,1),(1,2)

在指定的张量x中按照index0索引选择数据，得到z

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