一、函数介绍
scatter_(input, dim, index, src):将src中数据根据index中的索引按照dim的方向填进input。可以理解成放置元素或者修改元素
- dim:沿着哪个维度进行索引
- index:用来 scatter 的元素索引
- src:用来 scatter 的源元素,可以是一个标量或一个张量
二、实现原理
-
x = torch.rand(
2,
5)
-
-
#tensor([[0.1940, 0.3340, 0.8184, 0.4269, 0.5945],
-
# [0.2078, 0.5978, 0.0074, 0.0943, 0.0266]])
-
-
torch.zeros(
3,
5).scatter_(
0, torch.LongTensor([[
0,
1,
2,
0,
0], [
2,
0,
0,
1,
2]]), x)
-
LongTensor的shape刚好与x的shape对应,也就是LongTensor每个index指定x中一个数据的填充位置。dim=0,表示按行填充,主要理解按行填充。eg. LongTensor中的第1行第0列的值为2,即索引index=2,表示在第2行(从0开始)进行填充,对应到zeros(3, 5)中就是位置(2,0)。所以此处要求zeros(3, 5)的列数要与x列数相同,LongTensor中的index最大值应与zeros(3, 5)行数相一致,示意图如下
最终可以得到填充完整的zeros
-
tensor(
[[0.1940, 0.5978, 0.0074, 0.4269, 0.5945],
-
[0.0000, 0.3340, 0.0000, 0.0943, 0.0000],
-
[0.2078, 0.0000, 0.8184, 0.0000, 0.0266]])