torch.autograd.profiler.profile()で逆伝搬のボトルネックを探す(メモ)

記事というほどのものになっていないので、メモとして公開します。
検証が不足しているので、間違っているところがあったら教えていただけると嬉しいです。🙇‍♂️

---

以前こんなことを書いていました。遅くなってごめんなさい。。

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">Pytorchでbackwardが5分とかかかった時に使ったtorch.autograd.profilerとか便利だと思うので記事化したい。どこのbackwardで時間かかってるのか可視化できた。| <a href="https://t.co/efFK54Vw96">https://t.co/efFK54Vw96</a></p>&mdash; m3yrin😷 (@m3yrin) <a href="https://twitter.com/m3yrin/status/1123599620683120642?ref_src=twsrc%5Etfw">May 1, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

## お気持ち
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:`とかをおもむろに追加します。

たとえば

```python
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()`が終わらないとするとき、こんな感じにしてみます。

```python
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()`で結果をみます。

```python
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の形式でファイルに書き込めます。
ファイルパスは適宜変えてください。

```python
prof.export_chrome_trace("./trace.json")
```

##### Chrome tracing tools立ち上げ

Google Chromeのアドレスバーに`chrome://tracing`を入力すると立ち上がります(こんなのあったんだ...👀)

![1.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/329161/fd5048b8-c08e-1071-0a22-c5136f8899e3.png)


##### Loadのボタンからファイルを読み込こむ

先程書き出したファイルをLoadボタンから読み込みます

![1-2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/329161/b275db28-d201-f6dc-7fc9-9fcc8074441c.png)

すると、以下のようにBackwardの処理が時系列順に表示されます。

![2.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/329161/8c84504c-6c1a-f8de-dfcb-83e550c26b1a.png)


##### 4 探す
このままでは細かすぎてわからないので、カーソルを拡大縮小に切り替えてから拡大すると、それぞれの処理の名称が見えます。
書いたコードとこれを照らし合わせながら、時間がかかっているところを探します。

![3.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/329161/46b32033-d35d-abae-93ce-ee93db4e8c4c.png)

僕が以前失敗したときは、以下のようにEmbeddingの後ろにやたらと長い処理(緑色の長いやつ)がみえたので、時間がかかっている箇所が特定できました。

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">ちなみにですが、torch.autograd.profilerの結果は以下のように可視化できて、順番にみていくとEmbeddingの直後にボトルネック(緑の長いやつ)があったので、犯人こいつじゃん、と特定できました。 <a href="https://t.co/jRsD2Jcbvn">pic.twitter.com/jRsD2Jcbvn</a></p>&mdash; m3yrin😷 (@m3yrin) <a href="https://twitter.com/m3yrin/status/1123770278549557249?ref_src=twsrc%5Etfw">May 2, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

上記のように、一つのブロックで長い時間を消費している場合には、この可視化で発見しやすいです。
大きなブロックが見えない場合には、変なミスをしているのではなく、扱っているデータがそもそも重いなどが考えられます。

## おわりに

`torch.autograd.profiler`には他にも機能があって、もっといい方法がありそうな気がします。
https://pytorch.org/docs/stable/autograd.html#profiler

以上、参考になれば幸いです。