[論文レビュー] Profile-guided memory optimization for deep neural networks

Taro Sekiyama, Takashi Imamichi+2 more|arXiv (Cornell University)|Apr 26, 2018

Neural Networks and Applications被引用数 9

ひとこと要約

この論文は、1回のフォワードパスのプロファイリングから得られるヒューリスティックに基づくメモリ割り当て戦略を用いて、ピークメモリ使用量を最大49.5％まで削減する、プロファイル駆動型メモリ最適化を提案する。この手法により、動的メモリ割り当てのパフォーマンスオーバーヘッドが回避され、トレーニングおよびインフェンスの両方が最大4倍に加速され、プロファイリング後はランタイムのオーバーヘッドが一切ない。

ABSTRACT

Recent years have seen deep neural networks (DNNs) becoming wider and deeper to achieve better performance in many applications of AI. Such DNNs however require huge amounts of memory to store weights and intermediate results (e.g., activations, feature maps, etc.) in propagation. This requirement makes it difficult to run the DNNs on devices with limited, hard-to-extend memory, degrades the running time performance, and restricts the design of network models. We address this challenge by developing a novel profile-guided memory optimization to efficiently and quickly allocate memory blocks during the propagation in DNNs. The optimization utilizes a simple and fast heuristic algorithm based on the two-dimensional rectangle packing problem. Experimenting with well-known neural network models, we confirm that our method not only reduces the memory consumption by up to $49.5\%$ but also accelerates training and inference by up to a factor of four thanks to the rapidity of the memory allocation and the ability to use larger mini-batch sizes.

研究の動機と目的

• GPU やエッジシステムなどのメモリ制限のあるデバイスにおいて、深層ニューラルネットワーク（DNN）の高いメモリ消費量を軽減すること。
• モデルの計算を変更せずにピークメモリ使用量を削減し、より大きなミニバッチサイズと高いGPU利用度を実現すること。
• 深層学習フレームワークにおける動的メモリ割り当てのパフォーマンスオーバーヘッドを回避する、高速なプロファイル駆動型メモリ割り当て技術を開発すること。
• 可変長入力を有するモデル（例：RNN）に対しても、効率的なメモリ管理を可能にすること。

提案手法

• DNNの1回のフォワードパス中にメモリ使用量をプロファイリングし、『ホット』な計算パスのメモリアクセスパターンを捉える。
• ピークメモリ使用量を最小化するために、メモリ割り当て問題を二次元の長方形パッキング問題として定式化する。
• 近似的に最適なメモリブロック割り当てを効率的に生成できる、シンプルで高速なヒューリスティックアルゴリズムを開発する。
• プロファイルから事前に計算済みのメモリアドレスをキャッシュすることで、以降の実行で定数時間のメモリアクセスを実現する。
• Chainer深層学習フレームワークに統合され、他のフレームワークとも互換性がある。
• 可変長シーケンス（例：RNN）の文脈では、不要なブロックによるメモリブloatを避けるために、必要に応じて再計算を行う。

実験結果

リサーチクエスチョン

• RQ1 プロファイル駆動型メモリ割り当ては、ランタイムのオーバーヘッドなしに、深層ニューラルネットワークのピークメモリ使用量を顕著に削減できるか？
• RQ2 2次元の長方形パッキング問題に対して、高速なヒューリスティックベースのアプローチがDNNのメモリ消費量を最小化するのにどの程度有効か？
• RQ3 どの程度のメモリ削減が、トレーニングにおけるより大きなミニバッチサイズとGPU利用度の向上を可能にするか？
• RQ4 プロファイル計算のコストを考慮しても、インフェンス速度が維持または向上するか？
• RQ5 RNN のような可変長入力を持つモデルにおいて、この手法はどの程度のパフォーマンスを示すか？

主な発見

• ResNet-50 と Inception-ResNet では、ピークメモリ使用量を最大49.5％まで削減。推論時には、GoogLeNet で12.6％、ResNet-50 で10.0％の削減を達成。
• より大きなミニバッチサイズに起因する高速なメモリ割り当てとGPU利用度の向上により、トレーニング速度が最大4倍に加速。
• seq2seq モデルでは、メモリ割り当てコストの低減とより良いメモリ再利用のおかげで、インフェンスが最大23.8％高速化。
• ヒューリスティックアルゴリズムは実用的に十分に高速であり、ほとんどのモデルで1秒未満の実行時間で完了。
• RNN におけるプールベースのメモリ管理よりも優れており、メモリプールのフラグメンテーションを回避し、メモリの解放・再割り当ての必要性を低減。
• 128および256のミニバッチサイズにおいて、最適化版では元のChainerよりもトレーニングが高速である。これは、メモリ割り当てのオーバーヘッドが削減されたため。