PyTorchのPERFORMANCE TUNING GUIDEの効果を確認してみる その1 「parameter.grad = Noneを使う」

PyTorchには「PERFORMANCE TUNING GUIDE」という学習を速くするためのテクニック集があります。ただ、このドキュメントでは個々のテクニックでどれくらい速くなるのかまでは書いていません。このため、気になるものをピックアップして、具体的なコードとそのときの効果を実際に測って示そうと思います。

今回は1回目ということで、簡単にできる 「Use parameter.grad = None instead of model.zero_grad() or optimizer.zero_grad()」 の効果を確認してみます。

model.zero_grad() や optimizer.zero_grad() の代わりにparameter.grad = Noneを使うとは？

このテクニックがどういうことかを実際のコードを使って示します。PyTorchの学習コードは以下のようなループを繰り返すかと思います。

for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    output = model(data)
    loss = criterion(output, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

この中で、gradの初期化をしている optimizer.zero_grad() を以下のように置き換えるテクニックです。

for param in model.parameters():
    param.grad = None

これにより余計な初期化の処理などを省くことができ、学習が速くなると書かれています。また、PyTorchの1.7からは以下のように引数に set_to_none=True とすることで同様の効果が得られるようになりました。

optimizer.zero_grad(set_to_none=True)

実際に効果を測定してみる

さて、parameter.grad = None を使うことでどれくらい速くなるのか？というのが気になったので軽く測定してみます。実行環境は以下の通りです。

• 実行環境：Google Colab
• GPU: Tesla T4
• PyTorch: 1.8.1
• torchvision: 0.9.1
• 使用モデル: ResNet50

また、計算時間はoptimizer.zero_grad()のgradの初期化の部分とイテレーション全体の2種類を測定します。また、測定の際、CUDA側の計算時間を測定するためにtorch.cuda.Eventを利用しました。

より詳細に知りたい方は測定コードはこちらに上げておきましたので、そちらを確認してください。

https://github.com/shu65/pyorch_performance_tuning_guide_examples/blob/main/Use_parameter_grad_%3D_None.ipynb

測定結果は以下の通りです。

gradの初期化のところはCUDA側の処理がなくなっているので、CUDA側の処理時間はほぼ0になりました。ただ、元々、イテレーション全体でみるとほとんど時間がかかっていないので、今回のテクニックを使ってもイテレーション全体でみると何も変わらないどころか少し遅くなるという結果でした。私自身、普段ResNet50を使っているときに、gradの初期化のところに時間がかかっているというのはほとんど気になったことがないので、順当な結果かなと思います。ただし、ResNet50以外はイテレーション全体にかかる時間とgradの初期化にかかる時間の比が違うと思われるので、効果の大きさは変わると思われますのであくまで参考程度の結果だと思ってみてください。

終わりに

今回はPyTorchの「PERFORMANCE TUNING GUIDE」の中の「Use parameter.grad = None instead of model.zero_grad() or optimizer.zero_grad()」 の効果を確認してみました。今後も他のテクニックの効果も測定していこうと思います。