PyTorchでGPUの計算時間を正しく計測する

今回の記事ではPyTorchでGPUで実行した関数の計算時間を正しく測定する方法とその後に詳しい説明をしていきます。

はじめに

仕事がらPyTorchで高速な学習方法をいろいろ調べることがよくあります。
その際、blog記事などで、Pythonの time() を利用して計算時間を測定して「こんなに速くなりました！」という紹介記事を見かけることがあります。ただ、そこに載っているコードがGPU用の測定方法を用いていないため、正しく測定できていなくて数値が参考にならないということが本当によくあります。

せっかく、いいまとめなのにもったいない・・・ということが多いため、この記事では少しでもそういうものが減ってくれればと思い、PyTorchのGPUの処理の正しい計算時間測定方法についてまとめました。ちなみに、profilerを使ったやり方は別の記事にしようとかと思うので今回は言及しません。

サンプルコードはこちらにありますので、わからないところがあれば実際に動かして確かめてみてください。

https://github.com/shu65/pytorch-cuda-time-measurement/blob/main/Pytorch_GPU_Time_Measurement.ipynb

具体的なGPUの計算時間の測定方法

今回は torch.cuda.synchronize() と torch.cuda.Event を利用した2種類の方法を紹介します。

torch.cuda.synchronize() を利用した方法

torch.cuda.synchronize() を利用する場合は以下のように time()の前に torch.cuda.synchronize() を実行するようにします。

torch.cuda.synchronize()
start = time.time()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)  
# 測定したい部分終了
torch.cuda.synchronize()
elapsed_time = time.time() - start

print(elapsed_time, 'sec.')

torch.cuda.Eventを利用した方法

torch.cuda.Event を利用する場合は開始用と終了用のtorch.cuda.Event を作り、測定したい関数の前後でrecord()を呼びます。その後、GPUの処理が終わるまで待つために torch.cuda.synchronize() を呼び、elapsed_time() で計算時間を取得するという流れになります。

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)  
# 測定したい部分終了
end.record()
torch.cuda.synchronize()
elapsed_time = start.elapsed_time(end)

print(elapsed_time / 1000, 'sec.')

詳しい説明

まず、よくある間違えがどんなコードか？を説明します。よくある間違えているコードのサンプルとしては以下の通りです。

start = time.time()
# 測定したい部分開始
with torch.no_grad():
  out = model_gpu(input_batch_gpu)
# 測定したい部分終了
elapsed_time = time.time() - start

print(elapsed_time, 'sec.')

このように普通に time() を呼ぶという間違いが多い印象です。CPUではこれで問題ないのですが、GPUを使った場合はこれでは正しく計算時間が測定できていません。その理由はCPUとGPUの処理が非同期で行われているからです。

CPUとGPUの処理が非同期とは？

PyTorchではGPUの処理を実行する際にCUDAを利用しています。このCUDAではGPUで処理する関数をkernel関数、または単にkernelと呼びます。このkernel関数はCPUとGPUのリソースを最大限に活用できるように、基本的にはCPUがkernel関数の実行を依頼するGPUのタスクキューに積むところまで行い、kernel関数の処理が終わるのを待たないで返ってくるということをしています。

この結果、GPUが処理している最中もCPUが別の処理を実行でき、計算リソースを有効活用することができます。

GPUを活用するという意味ではこの非同期処理という仕組みは非常に有用なのですが、GPUで行っている処理の計算時間を測定する場合には注意が必要になります。というのも、kernel関数を呼んで返ってきたタイミングではGPUの処理が終わってないためです。このため、処理の時間を測定する場合はCPUとGPUとの同期をしたり、CUDA Eventなどの特別な方法で測定する必要があります。

正しく測定してない場合と正しく測定した場合どれくらい差がでるのか？

間違っていたらどれくらいひどいことになるか？を実感してもらうために、Google ColabでResNet50というモデルの推論を行ったときの結果を紹介します。コードはこちらになります

https://github.com/shu65/pytorch-cuda-time-measurement/blob/main/Pytorch_GPU_Time_Measurement.ipynb

実行環境は以下の通りです。

• GPU: T4
• CUDA: 10.1
• PyTorch: 1.8.0
• torchvision: 0.9.0

また、ResNet50の入力バッチサイズは128として実行します。

この条件で実行した結果は以下の通りです

	時間 (sec.)
CPU	18.83
GPU (間違った測定方法の場合)	0.01
GPU (torch.cuda.synchronize()利用時)	0.32
GPU (torch.cuda.Event利用時)	0.32

結果として、間違った測定方法だとCPUとGPUを比較すると「1883倍速くなりました！」という主張をしてしまうことになります。ちなみにGPUで1000倍なんて数字が出てきたら確実にどこか間違えています。実際、今回のケースでは本当は「約59倍速くなりました！」というのが正しい結果になります。

torch.cuda.synchronize()とtorch.cuda.Eventを使った場合の違い

今回torch.cuda.synchronize()とtorch.cuda.Event の2種類を紹介しました。場合によっては使い分けをしたほうがいいのでこの二つの違いを説明していきます。

torch.cuda.synchronize() を利用した場合、簡単なので測定しやすいのでいいので、ぱっと測定したい場合はこちらの方法が楽でよいかと思います。ただ、こちらの方法はkernel関数の発行と測定終了のtorch.cuda.synchronize() の終了までの時間も含むことになります。kernel関数の発行もtorch.cuda.synchronize() も時間としては十分短いことが多いので、ほとんどの場合は無視できると思います。ただ、常時監視する目的で測定する際には、torch.cuda.synchronize() を測定したい部分の終了時に呼ぶため、CPUの処理がtorch.cuda.synchronize()のところで止まってしまうのでオーバーヘッドが大きすぎるという問題があります。
一方、torch.cuda.Event を利用した場合はelapsed_time() を呼ぶ直前に何等かの方法でCPUと同期すればいいので、学習のイテレーションの最後に同期するなど工夫することができ、この結果、オーバーヘッドを小さくすることができます。このため、常時監視する目的で測定する際はtorch.cuda.Event の利用をお勧めします。

終わりに

今回はPyTorchのGPUの計算時間を正しく計測する方法について紹介しました。この記事でGPUの計算時間の測定方法を間違えておかしなことを主張する記事が少しでも減ってくれれば幸いです。
profilerについても今度調べて記事にできればと思っています。