TL;DR
GPUでNLPする時のCUDA out of memoryを回避する方法を地味なものからナウいものまでまとめてみた
自己紹介
- 都内のしがない博士院生
- NLPer
- PyTorchユーザー
- VAEが好き
CUDA out of memory とは
GPUのメモリに大きなサイズのテンソルを乗せすぎて容量を超えてしまうことで発生するエラー。
NLPerに限らず、多くのエンジニア・リサーチャーを苦しめていることと（勝手に）思っています。
この記事ではそんな方達のために僕が今までに見つけた回避術をいくつか紹介したいです。
他にもこんなの知っているよ、とか、ちょっとその方法はどうなの、という方がいたらコメントしていただけると勉強になるので嬉しいです。
以下、padding, embedding, RNN, seq2seqなどは前提知識として話を進めます。
回避術
1. 最大系列長を指定しよう
NLPerが扱うデータは自然言語なので可変長。
そしてデータの中には大抵外れ値が紛れ込んでいて、とても短いものや、長いものもあります。
特に長いものをそのままテンソルにしてGPUに載せると一瞬でout of memoryになることがあります。
以前、最大系列長を指定しておらず、batchごとの最大長さに合わせて長さを変えるということをしていましたが、
実際に手元で実験しているときに、スペース10万個からなる単語がbatchに紛れており(原因は不明)、そのbatchになると必ずout of memoryになるケースがありました。
文を扱うときは単語の総数が系列長に相当するが、このように単語を文字に分解してCNNに入力するときは、単語の長さも肝になってきます。
事前に最大文長、単語長を調べる等の簡単な統計量の調査くらいはしてもいいかもしれないです。
（よく見るとpytorch公式でも同様のアドバイスがされていた）
余談ですが、文字レベルCNNを使うときは、最低限CNNのfilterの幅だけは長さが必要になるので、（単語の）最小系列長を指定する必要もあることに触れておきます。
2. Embedding layerだけCPUに載せる
どうしてもGPUに乗らなくて、batch size 2じゃないと計算できない・・・ということがありました。
その時の裏技として、embedding layerだけCPUに乗せ、embedded vectorからGPUに載せるというものがあります。
単語数万x埋め込み次元300となるとembedding layerがかなりの容量を占めるので、embeddingだけ目を瞑ってCPUでやるのはなかなか効果的でした。
3. Tie embedding
RNN言語モデルでは、まず単語をembeddingして、なんやかんやした後、最終的に線形変換後、softmaxして語彙次元のテンソルにする、を各stepでします。
この時に、embeddingのパラメータと線形変換のパラメータをshareすることで省メモリになります。
具体的にはこの論文を参照されたい。
TYING WORD VECTORS AND WORD CLASSIFIERS: A LOSS FRAMEWORK FOR LANGUAGE MODELING (ICLR2017)
4. Share embedding
3と類似手法になるが、seq2seqモデルを使う時に、encoderとdecoderの入力側のembedding layerに同じものを使う手法。
3, 4については
https://github.com/howardyclo/pytorch-seq2seq-example
に詳細な実装が公開されています。（僕が書いたものではないです）
実際使ってみた上での注意点として、メモリ上でパラメータはshareしていてもモデルを保存するときは別々のものとして扱われるっぽいので、省ストレージするにはさらに工夫が必要。
5. 不要になったテンソルはこまめにdelする
例えば、batch処理が終わってloss.backward()した後は、次のbatch処理を始める前にdel lossした方がいい（かもしれない）(気がする)(気持ち的に)
6. Vocabulary selection
ラボの某先輩が教えてくれた論文で、テキスト分類するときは、実は1万も単語はいらない、なんなら1000単語でも十分な精度がでる話。
How Large a Vocabulary Does Text Classification Need? A Variational Approach to Vocabulary Selection (NAACL2019)
論文ではvariationalに単語を落としているけど、実用的には手動で設定してもいい気はする。
タスクの難易度によって必要な語彙数が違うのか気になる。
※6は実際に試したことはないです
最後に
advent calendarの投稿遅れてしまいましてすみません。
冒頭にも書きましたがVAEが好きです。
一緒に議論・論文読み会などできる方を探しているので、twitterでもメールでもしてくださるととても嬉しいです。 #VAE好きな人と繋がりたい
よろしくお願いします。