記事というほどのものになっていないので、メモとして公開します。
検証が不足しているので、間違っているところがあったら教えていただけると嬉しいです。🙇♂️
以前こんなことを書いていました。遅くなってごめんなさい。。
Pytorchでbackwardが5分とかかかった時に使ったtorch.autograd.profilerとか便利だと思うので記事化したい。どこのbackwardで時間かかってるのか可視化できた。| https://t.co/efFK54Vw96
— m3yrin😷 (@m3yrin) May 1, 2019
お気持ち
PyTorchで論文の再現実装とかに挑戦したりすると、loss.backward()の一行が遅すぎて悲しくなる時がよくあります。(僕はあります。。)
Backwardでエラーにはならないけど非常に時間がかかる場合、処理のどの箇所で時間がかかってるのかを探すのは手間だったりします。
Backwardでのボトルネックを探す一つの手がPyTorchのProfilerです。
https://pytorch.org/docs/stable/autograd.html#torch.autograd.profiler.profile
このメモでは、torch.autograd.profiler.profile()を使って、逆伝搬にかかる時間を可視化する手順を書きます。
やりかた
コードはPyTorchのTutorialsを例に使います。
NLP FROM SCRATCH: TRANSLATION WITH A SEQUENCE TO SEQUENCE NETWORK AND ATTENTION
https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html
また実行環境はJupyter Notebookを仮定しています。
手元の環境はPython 3.7.5, PyTorch 1.3.1です。
1.train等の処理をwith構文で囲む
Backwardの処理を含めるようにwith torch.autograd.profiler.profile() as prof:とかをおもむろに追加します。
たとえば
hidden_size = 256
encoder1 = EncoderRNN(input_lang.n_words, hidden_size).to(device)
attn_decoder1 = AttnDecoderRNN(hidden_size, output_lang.n_words, dropout_p=0.1).to(device)
trainIters(encoder1, attn_decoder1, 75000, print_every=5000)
でtrainIters()が終わらないとするとき、こんな感じにしてみます。
hidden_size = 256
encoder1 = EncoderRNN(input_lang.n_words, hidden_size).to(device)
attn_decoder1 = AttnDecoderRNN(hidden_size, output_lang.n_words, dropout_p=0.1).to(device)
with torch.autograd.profiler.profile(use_cuda=True) as prof:
trainIters(encoder1, attn_decoder1, 1, print_every=1)
イテレーションの回数とかは一応1とかにしました。手元の環境ではGPUを使っているのでuse_cuda=Trueにしてますが、cpuのときはFalseのほうが良いでしょう。また、with句の範囲は、もう少し狭くてもいいかと思います。
2. 結果をプリントする
2.1 key_averages()を使う
Docsだとこれが載ってます。
上の処理を実行してから、prof.key_averages()で結果をみます。
print(prof.key_averages().table(sort_by="self_cpu_time_total"))
を実行すると、以下のようになります。
(表示のため一部の列のみを表示しています)
------------------------------------------ --------------- --------------- ---------------
Name Self CPU total % Self CPU total CPU total %
------------------------------------------ --------------- --------------- ---------------
_cudnn_rnn_backward 16.41% 8.496ms 16.41%
add 12.71% 6.580ms 12.71%
_cudnn_rnn 8.75% 4.528ms 8.75%
SelectBackward 4.63% 2.395ms 12.14%
zero_ 3.69% 1.909ms 3.69%
mm 3.57% 1.845ms 3.57%
cat 3.52% 1.823ms 3.52%
addmm 3.45% 1.785ms 3.45%
empty 3.05% 1.576ms 3.05%
embedding_dense_backward 2.74% 1.421ms 2.74%
view 2.51% 1.297ms 2.51%
select 2.33% 1.208ms 2.33%
bmm 2.14% 1.105ms 2.14%
embedding 2.00% 1.035ms 2.00%
zeros 1.95% 1.007ms 7.00%
sum 1.91% 990.300us 1.91%
add_ 1.61% 832.000us 1.61%
_softmax_backward_data 1.36% 704.500us 1.36%
gru 1.29% 669.900us 11.05%
t 1.24% 644.300us 1.24%
struct torch::autograd::CopySlices 1.24% 642.700us 3.73%
relu 1.20% 620.900us 1.20%
実行時間順に処理の名称が表示されます。
これでボトルネックが分かればいいのですが、いまいちピンときません..
2.2 Chrome tracing toolsを使う
ファイルへの書き出し
export_chrome_trace()というメソッドがあったので、これを使ってみます。
以下のようなコードで、EventListをChrome tracing toolsの形式でファイルに書き込めます。
ファイルパスは適宜変えてください。
prof.export_chrome_trace("./trace.json")
Chrome tracing tools立ち上げ
Google Chromeのアドレスバーにchrome://tracingを入力すると立ち上がります(こんなのあったんだ...👀)
Loadのボタンからファイルを読み込こむ
先程書き出したファイルをLoadボタンから読み込みます
すると、以下のようにBackwardの処理が時系列順に表示されます。
4 探す
このままでは細かすぎてわからないので、カーソルを拡大縮小に切り替えてから拡大すると、それぞれの処理の名称が見えます。
書いたコードとこれを照らし合わせながら、時間がかかっているところを探します。
僕が以前失敗したときは、以下のようにEmbeddingの後ろにやたらと長い処理(緑色の長いやつ)がみえたので、時間がかかっている箇所が特定できました。
ちなみにですが、torch.autograd.profilerの結果は以下のように可視化できて、順番にみていくとEmbeddingの直後にボトルネック(緑の長いやつ)があったので、犯人こいつじゃん、と特定できました。 pic.twitter.com/jRsD2Jcbvn
— m3yrin😷 (@m3yrin) May 2, 2019
上記のように、一つのブロックで長い時間を消費している場合には、この可視化で発見しやすいです。
大きなブロックが見えない場合には、変なミスをしているのではなく、扱っているデータがそもそも重いなどが考えられます。
おわりに
torch.autograd.profilerには他にも機能があって、もっといい方法がありそうな気がします。
https://pytorch.org/docs/stable/autograd.html#profiler
以上、参考になれば幸いです。