PyTorch: requires_grad = False と optimizerにparamを与えないことの違い

疑問

つくりながら学ぶ! PyTorchによる発展ディープラーニングの転移学習を実装しているときに出てきた疑問。

モデルのパラメータにrequires_grad = Falseを設定することと、optimizerにparamを与えないことは動作的に異なるのだろうか。

書籍の中で、モデルのもつパラメータのうち、更新したいパラメータにはparam.requires_grad = Trueを、それ以外にはFalseを与えている。

# 転移学習で学習させるパラメータを、変数params_to_updateに格納する
params_to_update = []

# 学習させるパラメータ名
update_param_names = ["classifier.6.weight", "classifier.6.bias"]

# 学習させるパラメータ以外は勾配計算をなくし、変化しないように設定
for name, param in net.named_parameters():
    if name in update_param_names:
        param.requires_grad = True
        params_to_update.append(param)
        print(name)
    else:
        param.requires_grad = False

その後、optimizerにparam.requires_grad = Trueのパラメータのみ設定。

# 最適化手法の設定
optimizer = optim.SGD(params=params_to_update, lr=0.001, momentum=0.9)

勾配計算とその更新において、これらの操作がどう影響するのだろうか。

結論

PyTorchのDiscussionを参照した。

param.requires_grad = Falseをパラメータに与えた場合、勾配計算の対象外となり、その結果、勾配も更新されない。 このDiscussionで紹介されている以下のコードの通り、もしgradのないパラメータはp.gradにNoneを持ち、それらのupdate処理はskipされる。

                if p.grad is None:
                    continue

一方、optimizerのparamsに設定しない場合、勾配計算の更新の対象外となるものの、勾配の計算自体は行われる。

検証

転移学習のコードを改変して、以下の2パターンで速度を検証してみた。

1. 更新したいパラメータについてのみrequires_grad=True かつ optimizerには更新したいパラメータのみ指定
2. 全てのパラメータにおいてrequires_grad=True かつ optimizerには更新したいパラメータのみ指定
3. 更新したいパラメータについてのみrequires_grad=True かつ optimizerには全てのパラメータを指定

結果は、

1. 205秒
2. 458秒
3. 216秒

であった。

2においては、不必要に勾配計算している分遅いことが、また1と3については実質同じことであることが確認できた。

まとめ

単にoptimizerのparamsに設定しないだけだと、勾配の計算自体は全てのパラメータについて行われてしまう。

今回の転移学習では最後の層のパラメータのみ更新対象としたので、それ以降の勾配が計算されようがされまいが、時間は不要にかかるものの、結果に影響は無かった。

しかし、（ありうるケースなのかは不明だが）中間の層のパラメータのみ更新対象とした場合、その層から出力層までのパラメータ計算の影響を受けてしまうのでよろしくない。

更新不要なパラメータについては勾配計算もされないよう、更新したいパラメータについてのみrequires_grad=Trueを、それ以外はFalseを設定する必要がある。