PyTorch関数名の末尾アンダーバー`_`の意味

in-place処理であることを意味します。

はじめに
PyTorch関数名の末尾アンダーバー_
- 意味
- 具体例
おわりに
参考

はじめに

PyTorchを触っていると、しばしばアンダーバー_を接尾語とする関数を見かけます。

kaiming_normal_
add_
etc...

初見で意味を知らなかったので備忘録メモです。

PyTorch関数名の末尾アンダーバー`_`

意味

PyTorch関数名の末尾アンダーバー_は、in-place処理を意味します。

in-place処理とは、元のデータを直接変更する処理のことです。

In-place operations Operations that have a _ suffix are in-place.
For example: x.copy_(y), x.t_(), will change x.

Ref. Tensors — PyTorch Tutorials 2.4.0+cu121 documentation

in-place処理はメモリを節約することができますが、そこまでの履歴を書き換えることになるためautograd処理に問題が生じる可能性があるとのこと。基本的には非推奨です。

In-place operations save some memory, but can be problematic when computing derivatives because of an immediate loss of history. Hence, their use is discouraged.

Ref. Tensors — PyTorch Tutorials 2.4.0+cu121 documentation

具体例

ごく簡単な具体例として、addとadd_の違いを見ていきます。

まずはaddの例:

import torch


x = torch.tensor([1, 2, 3])
print("x: ", x)
y = x.add(1)
print("y: ", y)
print("x: ", x)

### 実行結果 ###
# x:  tensor([1, 2, 3])
# y:  tensor([2, 3, 4])
# x:  tensor([1, 2, 3])

add前後で、xには変化がありません。

一方、add_では、

x_ = torch.tensor([1, 2, 3])
print("x_: ", x_)
y_ = x_.add_(1)
print("y_: ", y_)
print("x_: ", x_)

### 実行結果 ###
# x_:  tensor([1, 2, 3])
# y_:  tensor([2, 3, 4])
# x_:  tensor([2, 3, 4])