PyTorch torch.scatter_add 関数


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

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)

出力結果は次のとおりです:

输入:
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 散布:
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)

注意:torch.scatter_add元の入力テンソルを変更せず、新しいテンソルを返します。複数のインデックスが同じ位置を指すことができ、値は累積されます。この関数はtorch.gatherの逆操作です。


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張機能