Cosnomi
Cosnomi

医療×IT / 医学生 / Web(React, Flask) / 機械学習(画像認識, Keras)

Twitter / GitHub / Keybase

GPUのメモリからCPUのメモリに逃がすテク「TFLMS」の論文を読んだ

Feb. 14, 2019機械学習論文

TFLMS という技術に関する論文を読んだので、軽くまとめます。(きっかけは第 1 回メディカル AI 学会で IBM の方が発表されていたからです。)

今後もこんな感じのを書いていきたいと思うのですが、あくまで自分の理解に基づいて書いているので、間違っている点などがあるかもしれません。実際に論文を利用する場合は必ず原著を参照してください。(また、コメントや Twitter などで指摘いただけるとありがたいです)

原著

Tung D. Le, Haruki Imai, Yasushi Negishi, Kiyokuni Kawachiya. 2018. TFLMS: Large Model Support in TensorFlow by Graph Rewriting (arxiv)

概要

大きい画像や 3D のタスクに deep learning を適用する際に、GPU メモリが制約となることがある。CPU のメモリに適切なタイミングで逃してやれば、その制約をなくせると考え、計算グラフに自動的にそのような処理を追加する tensorflow(と keras)のライブラリを開発したらしい。

こちらでコードが公開されている。

https://github.com/IBM/tensorflow-large-model-support

用語

swap-out: GPU のメモリから CPU のメモリにデータを逃がす

swap-in: CPU のメモリから GPU のメモリへデータを取り戻す

だから必ず、swap-out の後に swap-in が来る。

課題

swap-out, swap-in の操作はやはりロスが生じるので回数を減らしたい。また、GPU のメモリを節約するのが目的なので、出来るだけ早く swap out させたいし、できるだけ遅く swap in させたい

fusing

同じ tensor を複数回使うならswap out を複数回する必要はなくて、まとめてやれば良い。(swap-out operation を fuse する)

ある tensor が大きく、それを使うタイミングが近い場合、swap-in を fuse させることも考えられる。デメリットは、swap-in operation が早くなる(=その間、GPU メモリを専有する)こと。

タイミング

ここでは、swap-in のタイミングを考える。swap-in はできるだけ遅くしたいが、遅すぎると待機時間が生じる(overhead になる)。そこで、consuming operation(tensor を使う vertex)と swap-in vertex の control edge(swap-in を起動する辺みたいな)の topological order(vertex に付与される実行の順番みたいな)の差 k が一定範囲内に入るようにすることを考える。(この範囲はハイパーパラメータとなる)

どのようにしてこのような control edge を決めるかは 2 つの方法が提案されている。

direct-order strategy

k の下限から探索を始め、topological order がその値となる全ての vertices を取ってきて、consuming operation に reachable だったら、そこから control edge を生やす。

chain-rule strategy

こちらは、“chain-rule”から感じられるように具体的に neural network(NN)などを意識している。単にグラフとして考えているのではなく、もっと具体的に、forward と backward の存在する NN の computation graph として考え、その性質に着目して良い control edge を生成する。

NN の graph を考えると、forward の vertex から出る edge は次の layer に相当する forward の vertex へ行くものと、対応する backward の vertex に行くものがあるはず。ここで、forward に行くものについては、すぐ計算に使われていらなくなる。しかし、backward に使われるものについては、これ以降の layer の forward と backward の計算の後にようやく使われることになる。その間、GPU メモリを専有することとなり勿体無い。

だから、forward の vertex から breadth-first search をしていき、対応する backward の vertex について、k が範囲内にあれば、そこから control edge を生やす。

どうやら TFLMS では chain-rule strategy を採用している模様。

まとめ

実用性については、experiments を見る限りは上手く行っているようです。やっぱり、patch に分けず、そのままモデルに入れられれば、精度は上がりそうなので納得できます。


コメントフォームは設置していませんので、ご意見・ご感想などはTwitter(@cosnomi)などへお願いします。