pytorchで処理が重いとき on colaboratory

pytorch python

pytorchを使ったモデルで学習を行っている際、処理が重い場合があったためその調査/解決方法についてメモ

ちなみにgpuは基本的にgoogle colab上で使っているためその前提で

gpuが使えていない

処理が重い場合の一番の原因がこれ。その環境でgpuが利用可能か否かは以下でチェック

import torch
torch.cuda.is_available()

これでFalseが返ってくるならgpuは使えない状態になっているため環境を整える必要がある。

colab上であれば以下のケースが考えられる

  • ランタイムの選択がGPUになっていない
  • torchの間違ったversionがinstallされてしまった

私が遭遇した2つ目のケースの具体例として、torch関連のpackageをinstallする処理を実行したところ、dependencyの食い違いによりデフォルトで入っているtorchとは別のverisonがインストールされておかしくなっていたことがある。

以下を毎回colab起動時に実行していたのだが、ちょうど数日前に新しいversionがリリースされたことでtorchの依存versionが上がってしまいおかしくなっていた。なのでversionも指定しておくのがベター

!pip install torchaudio

いつも構築用のセルをまとめて実行して結果を見てない状態だったため、いつの間にか失敗してgpuが使えない状態になっていて(cpuでは普通に学習が動く状態)少しハマった。実は次に書いたプロファイルの確認をやってみて初めてgpuで動いてないことに気がついた、、

cpu上での処理が重い

gpuを使う状態になっていたとしてもcpu側で重い処理を行っていた場合は当然重くなる。この場合はプロファイルを取ってボトルネックになっている部分を探すしかない。 こんな感じで実行すればOK

!python -m cProfile -o trainer.prof trainer.py

sts = pstats.Stats('trainer.prof')
sts.strip_dirs().sort_stats('tottime').print_stats()

tottimeはおそらくtotal timeのことで、各メソッドでの合計処理時間(サブルーチンでの処理を除いた分)が計算される

例えばこんな出力になる

 Ordered by: internal time

 ncalls  tottime  percall  cumtime  percall filename:lineno(function)
  20836    9.360    0.000    9.361    0.000 {method 'to' of 'torch._C._TensorBase' objects}
  20800    1.822    0.000    1.822    0.000 {built-in method conv1d}
  20800    1.813    0.000    1.813    0.000 {built-in method conv2d}
  41600    1.293    0.000    1.293    0.000 {built-in method batch_norm}

ちなみにこれはcpu上で学習処理を実行したもの。そこまで大きいモデルでもbatch sizeでもないのに秒単位でconvに時間がかかっていることからわかる

プロファイルについてはこちらが参考になる
https://qiita.com/meshidenn/items/4dbde22d1e7a13a255bb