PyTorch torch.scatter_add 関数
torch.scatter_addは、PyTorch でソーステンソルの値を指定された位置に加算するための関数です。これはsrcの値を、indexで指定された位置に加算して、inputに反映します。
関数定義
torch.scatter_add(input, dim, index, src)
パラメータ:
input(Tensor): 入力テンソル。dim(int): 散布する次元。index(Tensor): インデックステンソル。src の値を input のどの位置に加算するかを指定します。src(Tensor): ソーステンソル。加算する値です。
戻り値:
torch.Tensor: 変更されたテンソルを返します。
使用例
例
import torch
# 入力テンソルを作成
input = torch.zeros(3, 5)
# インデックスとソースを作成
index = torch.tensor([[0, 1, 2, 0, 0],
[1, 2, 0, 1, 2],
[2, 0, 1, 2, 0]])
src = torch.tensor([[1, 1, 1, 1, 1],
[2, 2, 2, 2, 2],
[3, 3, 3, 3, 3]])
# dim=0 に沿って散布し累積
output = torch.scatter_add(input, dim=0, index=index, src=src)
print("入力:")
print(input)
print("nインデックス:")
print(index)
print("nソース:")
print(src)
print("n結果:")
print(output)
# 入力テンソルを作成
input = torch.zeros(3, 5)
# インデックスとソースを作成
index = torch.tensor([[0, 1, 2, 0, 0],
[1, 2, 0, 1, 2],
[2, 0, 1, 2, 0]])
src = torch.tensor([[1, 1, 1, 1, 1],
[2, 2, 2, 2, 2],
[3, 3, 3, 3, 3]])
# dim=0 に沿って散布し累積
output = torch.scatter_add(input, dim=0, index=index, src=src)
print("入力:")
print(input)
print("nインデックス:")
print(index)
print("nソース:")
print(src)
print("n結果:")
print(output)
出力結果は次のとおりです:
输入:
tensor([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]])
索引:
tensor([[0, 1, 2, 0, 0],
[1, 2, 0, 1, 2],
[2, 0, 1, 2, 0]])
源:
tensor([[1., 1., 1., 1., 1.],
[2., 2., 2., 2., 2.],
[3., 3., 3., 3., 3.]])
结果:
tensor([[4., 1., 2., 4., 4.],
[2., 1., 2., 2., 2.],
[3., 3., 1., 3., 3.]])
例
import torch
# dim=1 を使用
input = torch.zeros(3, 5)
index = torch.tensor([[0, 1, 2, 1, 0],
[1, 2, 0, 2, 1],
[0, 1, 1, 0, 2]])
src = torch.arange(1, 6).float()
output = torch.scatter_add(input, dim=1, index=index, src=src)
print("dim=1 に沿って散布:")
print(output)
# dim=1 を使用
input = torch.zeros(3, 5)
index = torch.tensor([[0, 1, 2, 1, 0],
[1, 2, 0, 2, 1],
[0, 1, 1, 0, 2]])
src = torch.arange(1, 6).float()
output = torch.scatter_add(input, dim=1, index=index, src=src)
print("dim=1 に沿って散布:")
print(output)
出力結果は次のとおりです:
沿 dim=1 散布:
tensor([[ 6., 3., 3., 0., 0.],
[ 3., 6., 3., 0., 0.],
[ 2., 4., 5., 0., 0.]])
例
import torch
# 複数の位置の値を集約する応用シナリオ
# 例:グラフニューラルネットワークで近傍ノードの特徴を累積する
# 4つのノードの初期特徴をシミュレートする
node_features = torch.zeros(4, 3)
# エッジの接続関係をシミュレートする(ソースノードがターゲットノードを指す)
edge_index = torch.tensor([0, 1, 2, 3, 0, 1]) # エッジのソースノード
edge_weights = torch.tensor([1.0, 2.0, 3.0, 1.5, 2.5, 0.5])
# 各エッジについて、ソースノード特徴の重み付き値を作成する
src_features = torch.randn(6, 3) * edge_weights.unsqueeze(1)
# 特徴をターゲットノードに累積する(ここでは簡略化している。実際にはエッジのターゲットノードに基づく必要がある)
target_nodes = torch.tensor([0, 0, 1, 1, 2, 3])
index = target_nodes
output = torch.scatter_add(node_features, 0, index.unsqueeze(1).expand_as(src_features), src_features)
print("ノード特徴の形状:", node_features.shape)
print("累積後の特徴:", output)
# 複数の位置の値を集約する応用シナリオ
# 例:グラフニューラルネットワークで近傍ノードの特徴を累積する
# 4つのノードの初期特徴をシミュレートする
node_features = torch.zeros(4, 3)
# エッジの接続関係をシミュレートする(ソースノードがターゲットノードを指す)
edge_index = torch.tensor([0, 1, 2, 3, 0, 1]) # エッジのソースノード
edge_weights = torch.tensor([1.0, 2.0, 3.0, 1.5, 2.5, 0.5])
# 各エッジについて、ソースノード特徴の重み付き値を作成する
src_features = torch.randn(6, 3) * edge_weights.unsqueeze(1)
# 特徴をターゲットノードに累積する(ここでは簡略化している。実際にはエッジのターゲットノードに基づく必要がある)
target_nodes = torch.tensor([0, 0, 1, 1, 2, 3])
index = target_nodes
output = torch.scatter_add(node_features, 0, index.unsqueeze(1).expand_as(src_features), src_features)
print("ノード特徴の形状:", node_features.shape)
print("累積後の特徴:", output)
注意:torch.scatter_add元の入力テンソルを変更せず、新しいテンソルを返します。複数のインデックスが同じ位置を指すことができ、値は累積されます。この関数はtorch.gatherの逆操作です。
その他の拡張機能
Pytorch torch リファレンスマニュアル