PyTorch笔记之scatter()函数

it2026-03-02  3

scatter(dim, index, src) 的参数有 3 个

dim:沿着哪个维度进行索引index:用来 scatter 的元素索引src:用来 scatter 的源元素,可以是一个标量或一个张量

直白点说就是在dim维度上以index将src中的数放到一个新的Tensor中。

首先说一下官网的计算规则:

self[index[i][j][k]][j][k] = src[i][j][k] # if dim == 0 self[i][index[i][j][k]][k] = src[i][j][k] # if dim == 1 self[i][j][index[i][j][k]] = src[i][j][k] # if dim == 2

从例子中可以看出,我们做的就是一个赋值操作,只不过不是直接self和src对应位置的赋值,而是根据index的值来在self的对应位置赋值。

另外需要注意,index中的值是不应大于该dim的通道数的。例如self是(3, 4, 5)的Tensor,我们选择的dim是2,则index的值必须小于5。

下面是两个例子:

1、one-hot编码(src是标量的情况下):

ind = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]) y = torch.zeros(3, 5).scatter_(0, ind, 1) print(y) ''' 输出: tensor([[1., 1., 1., 0., 0.], [1., 1., 1., 0., 0.], [0., 0., 0., 0., 0.]]) '''

详细解释:

我们可以将该函数看作两个过程:

首先进行坐标的映射,根据公式 self[index[i][j][k]][j] = src[i][j],所以其实这个坐标中j是不变的(绿色值),我们只需要把index对应位置的值填入作为新的i即可。表中红色的值是坐标。

然后在y的对应位置填入1即可。(坐标这个多画了一行,请忽略)

 

2、 src是张量的情况下: 

x = torch.rand(2, 5) ind = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]) y = torch.zeros(3, 5).scatter_(0, ind, x) print(x) print(y) ''' 输出: tensor([[0.1161, 0.8937, 0.7539, 0.9041, 0.5330], [0.7159, 0.1167, 0.0778, 0.9715, 0.7562]]) tensor([[0.1161, 0.1167, 0.0778, 0.9041, 0.5330], [0.0000, 0.8937, 0.0000, 0.9715, 0.0000], [0.7159, 0.0000, 0.7539, 0.0000, 0.7562]]) '''

详细解释:

与示例1类似,我们先得到对应的坐标映射,然后将src中(i, j)对应数据填到self的(index[i][j], j)中。

 

 

 

 

 

 

 

 

最新回复(0)