みーのぺーじ

みーが趣味でやっているPCやソフトウェアについて.Python, Javascript, Processing, Unityなど.

PyTorch の学習速度をなるべく速くする

PyTorch の学習に時間がかかるとイライラするので,利用可能な計算資源をなるべく効率的に使用するための工夫をまとめました.

以下は インテル Core i5 12400GeForce RTX 4080 を搭載したパソコンに Ubuntu をインストールした環境で検証しています.

GPU を使用していることを確認する

まず,nvidia-smi コマンドが正常に実行できるか確認します.これで GPU 使用率が取得できない場合は,環境がうまく構築できていません.

PyTorch を実行中に GPU 使用率を確認します.

0% ならば,torch.cuda.is_available()True であり,データとモデルを GPU に転送 (.to(device)) し忘れていないか再度確認します.

torch.Tensor.to — PyTorch 2.1 documentation

GPU 使用率が 0% より高い場合は,watch nvidia-smi コマンドを実行して,GPU 使用率を時系列で取得します.

GPU 使用率が低く変動する場合

GPU が効率的に使えていません.GPU メモリへのデータ転送待ちが発生している場合は,DataLoadernum_workers=1, pin_memory=True を指定することで改善する場合があります.

torch.utils.data — PyTorch 2.1 documentation

GPU 使用率が低く一定している場合

GPU の並列処理が効率的に動作していません.DataLoaderbatch_size を上げることで並列化可能な計算量が増えて GPU 使用率が上がり,学習速度が速くなる場合があります.しかし batch_size を大きくし過ぎるとミニバッチ学習のメリットが失われるかもしれません.

GPU 使用率が100%の場合

GPU の性能が不足しているので,許容できないほど学習速度が遅い場合は,より高性能な GPU に買い替えを検討します.

torch.compile を使用することで高速化する場合があるようですが,訓練の最初に追加の処理が必要になるので遅くなる場合もあります.

TorchScript を使用することで高速化する場合があるようですが,モデルが複雑な場合実行できないことがあります.

torch.amp で mixed precision を使用することで高速化する場合があるようですが,計算精度が落ちるので勾配消失やオーバーフローが発生しやすくなる可能性があります.

訓練速度を過度に最適化しないよう注意する

機械学習の目的はあくまでも汎化性能が高い予測モデルを構築することなので,訓練速度を上げることに情熱をかけすぎないように注意します.お茶を飲んだり,一晩寝たりすれば,その間に訓練が終わっているかもしれません.

とりあえず優先的に確認・検討すべき事項を分類してまとめました.