【初心者向け】PyTorchで特定のインデックス位置に値を加算する：`torch.Tensor.index_add()` 関数徹底解説

self テンサーは更新され、返される値となります。
indices と source の形状は、dim 次元を除いて一致する必要があります。
self テンサーの指定された次元 dim において、indices で指定されたインデックス位置に、source テンサーの値を alpha 倍して加算します。

構文

torch.Tensor.index_add_(self, dim, indices, source, alpha=1)

引数

alpha (float, optional): 加算する値の係数。デフォルトは 1
source (Tensor): 加算する値を格納したテンサー
indices (Tensor): 加算対象となるインデックスを格納したテンサー
dim (int): 加算対象の次元
self (Tensor): 更新対象のテンサー

戻り値

更新された self テンサー

詳細説明

indices テンサーは、self テンサーの dim 次元におけるインデックスを格納する必要があります。インデックスは整数型でなければならず、self.size(dim) より小さい値である必要があります。
source テンサーは、self テンサーと dim 次元を除いて同じ形状である必要があります。
alpha は、source テンサーの値を self テンサーに加算する際の係数です。デフォルトは 1 であり、source テンサーの値そのままが加算されます。

例

以下の例では、3 次元テンサー self において、1 次元目のインデックス 0 と 2 に、source テンサーの値をそれぞれ 2 倍して加算します。

import torch

self = torch.zeros(3, 4, 5)
indices = torch.tensor([0, 2])
source = torch.tensor([3, 4, 5])

self.index_add_(1, indices, source, alpha=2)
print(self)

このコードを実行すると、以下の出力が得られます。

tensor([[[3., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[0., 4., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[0., 0., 5., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]]])

PyTorch のバージョン 1.10 以降では、torch.scatter_add() 関数も提供されています。torch.scatter_add() 関数は、index_add() 関数と同様の機能を提供しますが、より柔軟な操作が可能になっています。
疎行列の操作によく用いられます。疎行列においては、非ゼロ要素のみを保持することでメモリ使用量を抑えることができます。index_add() 関数を利用することで、効率的に疎行列を更新することができます。
torch.Tensor.index_add() 関数は、inplace 操作です。つまり、self テンサー自身が更新され、新しいテンサーが返されるわけではありません。

3 次元テンサーへの値加算

この例では、3 次元テンサー self において、特定のインデックス位置に値を alpha 倍して加算します。

import torch

self = torch.zeros(3, 4, 5)
indices = torch.tensor([[0, 1], [2, 3]])
values = torch.tensor([3, 4, 5])
alpha = 2

self.index_add_(1, indices, values, alpha=alpha)
print(self)

tensor([[[3., 4., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[0., 0., 4., 8., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[0., 0., 0., 0., 5.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]]])

疎行列の更新

この例では、疎行列を torch.Tensor.index_add() 関数を使って更新します。

import torch
import scipy.sparse as sp

row_indices = torch.tensor([0, 1, 2])
col_indices = torch.tensor([2, 0, 1])
values = torch.tensor([3, 4, 5])
size = (3, 3)

A = sp.csr_matrix((values, (row_indices, col_indices)), shape=size)
B = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

A.data = A.data + B[row_indices, col_indices]
print(A.todense())

[[3. 0. 5.]
 [4. 5. 6.]
 [7. 8. 14.]]

バッチ処理

この例では、バッチ処理における torch.Tensor.index_add() 関数の使用方法を示します。

import torch

batch_size = 2
indices = torch.tensor([[0, 1], [2, 3]])
values = torch.tensor([[3, 4, 5], [6, 7, 8]])
alpha = 2

self = torch.zeros(batch_size, 4, 5)
for b in range(batch_size):
    self.index_add_(1, indices[b], values[b], alpha=alpha)

print(self)

tensor([[[3., 4., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[0., 0., 6., 12., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]]])

PyTorch の最新バージョンでは、より多くの機能が提供されている可能性があります。詳細は公式ドキュメントを参照してください。

手動ループ

最も単純な代替方法は、手動ループを使用してインデックスごとに値を更新することです。これは比較的単純な方法ですが、コードが冗長になり、計算速度が遅くなる可能性があります。

import torch

self = torch.zeros(3, 4, 5)
indices = torch.tensor([[0, 1], [2, 3]])
source = torch.tensor([3, 4, 5])
alpha = 2

for b in range(self.size(0)):
    for i in indices[b]:
        for j in range(self.size(2)):
            self[b, i, j] += alpha * source[j]

print(self)

利点

コードがわかりやすい

欠点

計算速度が遅い
コードが冗長になる

torch.scatter_add() 関数

PyTorch 1.10 以降では、torch.scatter_add() 関数が提供されています。この関数は index_add() 関数と同様の機能を提供しますが、より柔軟な操作が可能になっています。

import torch

self = torch.zeros(3, 4, 5)
indices = torch.tensor([[0, 1], [2, 3]])
source = torch.tensor([3, 4, 5])
dim = 1

self.scatter_add_(dim, indices, source, alpha=2)
print(self)

利点

コードが簡潔になる
index_add() 関数よりも柔軟性が高い

欠点

index_add() 関数よりも新しい機能なので、古いバージョンの PyTorch では利用できない

他のライブラリ

NumPy や SciPy などの他のライブラリを使用して、インデックス位置への値の更新を実行することもできます。これらのライブラリは、PyTorch よりも高速で効率的な場合があるため、パフォーマンスが重要な場合は検討する価値があります。

利点

PyTorchよりも高速で効率的な場合がある

欠点

コードが複雑になる場合がある
PyTorch との互換性が低い場合がある

カスタム関数

特定のニーズに合わせたカスタム関数を作成することもできます。これは、複雑な操作や、他のライブラリとの統合が必要な場合に役立ちます。

利点

他のライブラリとの統合が可能
特定のニーズに合わせた処理が可能

欠点

デバッグが難しい場合がある
コードの作成とメンテナンスに時間がかかる

最適な代替方法の選択

最適な代替方法は、状況によって異なります。以下の要素を考慮する必要があります。

互換性
柔軟性
コードの簡潔性
パフォーマンス