动手学习深度学习 | python函数学习

torch.gather/scatter_使用
说明：dim=0（按照竖直方向操作，即↓），dim=1（按照水平方向操作，即→）

gather是取的意思，scatter_是放的意思。

gatter

用法：output=torch.gather(input, dim, index)

a = torch.Tensor([[1,2],[3,4]])
b = torch.gather(a, 1, torch.LongTensor([[0,0],[1,0]]))
print(b) 

这里的“取”的意思是根据gather 函数中的参数从a中取出数据放入到变量b中， dim为 1，表示水平方向从 a 中取数据，a 中的数据对应的索引为[0,1][0,1]， 然后 根据取数据的索引LongTensor([[0,0],[1,0]])从a中对应的行取出数据，LongTensor([[0,0],[1,0]])表示，[0,0]表示从第一行取两次索引为0的数据1，[1,0]表示从 第二行分别取出索引为 1 和 0 的数据[4,3]。LongTensor 表示将取出的数据转化为Long型。所以 最后b为 tensor([[1.,1.],[4.,3.]])

scatter_

用法：output = torch.Tensor.scatter_(dim, index, src)

a = torch.rand(2, 5)
print(a)
b