torch.scatter

article/2025/11/11 10:04:15

本文目录

- 一、函数简介
- 二、二维举例
- 三、详解执行过程
- - 1. 第一步
  - 2. 第二步
  - 3. 第三步
  - 4. 问题

一、函数简介

torch.scatter(input, dim, index, src)

dim ([int]) – the axis along which to index
index (LongTensor) – the indices of elements to scatter, can be either empty or of the same dimensionality as src. When empty, the operation returns self unchanged.
src ([Tensor] or [float] – the source element(s) to scatter.
reduce ([str], optional) – reduction operation to apply, can be either 'add' or 'multiply'.

将src中的数据根据index中的索引按照dim的方向填入到input中

Writes all values from the tensor src into self at the indices specified in the index tensor. For each value in src, its output index is specified by its index in src for dimension != dim and by the corresponding value in index for dimension = dim.

看了上述的官方文档还是不理解，我们继续看官方的例子，这里官方只给了三维，我在这里又加入了二维，在这之前有一个规定

对任意维度d：index.size(d) <= src.size(d)
对d!=dim的维度：index.size(d) <= self.size(d)

二、二维举例

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

先上代码

torch.manual_seed(0)
x = torch.arange(0, 12).reshape(2, 6)
x= x.type(torch.float32)
print(x)
'''
tensor([[ 0.,  1.,  2.,  3.,  4.,  5.],[ 6.,  7.,  8.,  9., 10., 11.]])
'''index = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
print(index)
'''
tensor([[0, 1, 2, 0, 0],[2, 0, 0, 1, 2]])
'''y = torch.zeros(3, 6)
y = torch.scatter(y, 0, index, x)
print(y)
'''
tensor([[ 0.,  7.,  8.,  3.,  4.,  0.],[ 0.,  1.,  0.,  9.,  0.,  0.],[ 6.,  0.,  2.,  0., 10.,  0.]])
''''''in-place operation'''
yy = torch.zeros(3, 6)
yy.scatter_(0, index, x)
print(yy)
'''
tensor([[ 0.,  7.,  8.,  3.,  4.,  0.],[ 0.,  1.,  0.,  9.,  0.,  0.],[ 6.,  0.,  2.,  0., 10.,  0.]])
'''