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 ,