Pytorch 学习（5）：Pytorch中的 torch.gather/scatter_ 聚集/分散操作

 Pytorch 学习（5）：Pytorch中的 torch.gather/scatter_ 聚集/分散操作 

• torch.gather(input, dim, index, out=None) → Tensor 聚集操作
• torch.Tensor scatter_(dim, index, src) → Tensor 分散操作 

1.torch.gather(input, dim, index, out=None) → Tensor 聚集操作

沿给定轴 dim ,将输入索引张量 index 指定位置的值进行聚合，对一个 3 维张量,输出可以定义为:

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

如果 input 是 size 为 (x0,x1...,xi−1,xi,xi+1,...,xn−1)  且 dim = i 的 n 维张量,则 index 必须是具有 size 为 (x0,x1,...,xi−1,y,xi+1,...,xn−1)  的 n 维张量,其中 y >= 1