快捷導(dǎo)航

pytorch更新tensor中指定index位置的值scatter_add_問(wèn)題

更新時(shí)間：2023年06月14日 09:07:37 作者：騰陽(yáng)山泥若

這篇文章主要介紹了pytorch更新tensor中指定index位置的值scatter_add_問(wèn)題，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。如有錯(cuò)誤或未考慮完全的地方，望不吝賜教

使用scatter_add_更新tensor張量中指定index位置的值

例子

import torch
a = torch.zeros((3, 4))
print(a)
"""
tensor([[0., 0., 0., 0.],
? ? ? ? [0., 0., 0., 0.],
? ? ? ? [0., 0., 0., 0.]])
"""
b = torch.rand((2, 4))
print(b)
"""
tensor([[0.6293, 0.3050, 0.9608, 0.5577],
? ? ? ? [0.3469, 0.1025, 0.8185, 0.5085]])
"""
# 將a中第0行和第2行的值修改為b
a = a.scatter_add_(0, torch.tensor([[0, 0, 0], [2, 2, 2]]), b)
print(a)
"""
tensor([[0.6293, 0.3050, 0.9608, 0.0000],
? ? ? ? [0.0000, 0.0000, 0.0000, 0.0000],
? ? ? ? [0.3469, 0.1025, 0.8185, 0.0000]])
"""

torch_scatter.scatter_add、Tensor.scatter_add_ 、Tensor.scatter_、Tensor.scatter_add 、Tensor.scatter

torch_scatter.scatter_add

官方文檔：

torch_scatter.scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0)

Sums all values from the src tensor into out at the indices specified in the index tensor along a given axis dim. For each value in src, its output index is specified by its index in input for dimensions outside of dim and by the corresponding value in index for dimension dim. If multiple indices reference the same location, their contributions add.

看著挺疑惑的，自己試了一把：

src = torch.tensor([10, 20, 30, 40, 1, 2, 2, 2, 9])
index = torch.tensor([2, 1, 1, 1, 1, 1, 1, 1, 0])
out=scatter_add(src, index)
print(out)

輸出結(jié)果為：tensor([ 9, 97, 10])

說(shuō)白了就是：index就是out的下標(biāo)，將src所有和此下標(biāo)對(duì)應(yīng)的值加起來(lái)，就是out的值。

例如上面的例子：index中等于1的，對(duì)應(yīng)于src是【20, 30, 40, 1, 2, 2, 2】，將這些值加起來(lái)是97，于是，out[1]=97

同理：out[0]=src[8]=9 out[2]=src[0]=10

另一個(gè)函數(shù)

Tensor.scatter_add_

官方文檔：

scatter_add_(self, dim, index, other):

For a 3-D tensor, :attr:`self` is updated as::
? ? self[index[i][j][k]][j][k] += other[i][j][k] ?# if dim == 0
? ? self[i][index[i][j][k]][k] += other[i][j][k] ?# if dim == 1
? ? self[i][j][index[i][j][k]] += other[i][j][k] ?# if dim == 2

官方例子：

? ? ? ? ? ? >>> x = torch.rand(2, 5)
? ? ? ? ? ? >>> x
? ? ? ? ? ? tensor([[0.7404, 0.0427, 0.6480, 0.3806, 0.8328],
? ? ? ? ? ? ? ? ? ? [0.7953, 0.2009, 0.9154, 0.6782, 0.9620]])
? ? ? ? ? ? >>> torch.ones(3, 5).scatter_add_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
? ? ? ? ? ? tensor([[1.7404, 1.2009, 1.9154, 1.3806, 1.8328],
? ? ? ? ? ? ? ? ? ? [1.0000, 1.0427, 1.0000, 1.6782, 1.0000],
? ? ? ? ? ? ? ? ? ? [1.7953, 1.0000, 1.6480, 1.0000, 1.9620]])

以index來(lái)遍歷，就比較容易看懂。self中并不是每個(gè)值都要改變的。

以上面為例

index[0][0]=0 ?self[index[0][0]][0]=self[0][0] =self[0][0]+ x[0][0]=1 +0.7404=1.7404
index[0][1]=1 ?self[index[0][1]][1]=self[1][1] =self[1][1]+ x[0][1] =1 +0.0427 =1.0427

。。。

以此類推，將index遍歷一遍，就得到最終的結(jié)果

所以，self中需要改變的是index中列出的坐標(biāo)，其他的是不動(dòng)的。

Tensor.scatter_

scatter_(self, dim, index, src)

和Tensor.scatter_add_的區(qū)別是直接將src中的值填充到self中，不做相加

例子：

>>> x = torch.rand(2, 5)
? ? ? ? ? ? >>> x
? ? ? ? ? ? tensor([[ 0.3992, ?0.2908, ?0.9044, ?0.4850, ?0.6004],
? ? ? ? ? ? ? ? ? ? [ 0.5735, ?0.9006, ?0.6797, ?0.4152, ?0.1732]])
? ? ? ? ? ? >>> torch.zeros(3, 5).scatter_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
? ? ? ? ? ? tensor([[ 0.3992, ?0.9006, ?0.6797, ?0.4850, ?0.6004],
? ? ? ? ? ? ? ? ? ? [ 0.0000, ?0.2908, ?0.0000, ?0.4152, ?0.0000],
? ? ? ? ? ? ? ? ? ? [ 0.5735, ?0.0000, ?0.9044, ?0.0000, ?0.1732]])
? ? ? ? ? ? >>> z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
? ? ? ? ? ? >>> z
? ? ? ? ? ? tensor([[ 0.0000, ?0.0000, ?1.2300, ?0.0000],
? ? ? ? ? ? ? ? ? ? [ 0.0000, ?0.0000, ?0.0000, ?1.2300]])

另外，pytorch中還有

scatter_add和scatter函數(shù)，和上面兩個(gè)函數(shù)不同的是這個(gè)兩個(gè)函數(shù)不改變self，會(huì)返回結(jié)果值；上面兩個(gè)函數(shù)（scatter_add_和scatter_）是直接在原數(shù)據(jù)self上進(jìn)行修改