PyTorchのzero_grad()について

Question

PyTorchの解説記事を見ていると、

model.zero_grad()

optimizer.zero_grad()

という使われ方をしていました。
個人的には、モデルの勾配をリセットするのはマズいんじゃないかと思うのですが、どうなのでしょう。
かといって、オプティマイザーの勾配とは？という疑問もあり、どちらの方がより効率的なのか分かりません。

この２種類の違いや使い分けなど教えていただきたいです。

Answer 1

zero_grad() をする理由

Pytorch では、逆伝搬で勾配を計算した際に、以前に計算した勾配がある場合はそれに加算する形になっています。

今回の勾配 = 前回計算した勾配 + 今回計算した勾配

この仕様は RNN などのモデルでは有用なのですが、そうでない場合は前回計算した勾配を0で初期化する必要があるので、その場合は backward() を呼び出す前に zero_grad() を行う必要があります。

Optimizer.zero_grad() と nn.Module.zero_grad() の違い

torch.nn — PyTorch master documentation

Sets gradients of all model parameters to zero.
モデルのすべてのパラメータの勾配を0にする。

torch.optim — PyTorch master documentation

Clears the gradients of all optimized torch.Tensor s.
最適化対象のすべてのパラメータの勾配を0にする。

Optimizer.zero_grad() は最適化対象のすべてのテンソルの勾配を0で初期化します。
他方、nn.Module.zero_grad() はそのモジュールを構成するすべてのテンソルの勾配を初期化します。

なので、nn.Module オブジェクトのすべてのテンソルが最適化対象とした場合 (例: optimizer = optim.SGD(model.parameters()))、これら2つは同じ処理になります。

結論

backward() の前に Optimizer.zero_grad() を呼び出せばよい。