サブロウ丸

サブロウ丸

主にプログラミングと数学

サーベイ: Training Deep Nets with Sublinear Memory Cost

Chen, Tianqi, et al. "Training deep nets with sublinear memory cost." arXiv preprint arXiv:1604.06174 (2016).

@article{chen2016training,
  title={Training deep nets with sublinear memory cost},
  author={Chen, Tianqi and Xu, Bing and Zhang, Chiyuan and Guestrin, Carlos},
  journal={arXiv preprint arXiv:1604.06174},
  year={2016}
}

paper: https://arxiv.org/pdf/1604.06174

slide: https://www.cs.cmu.edu/~zhihaoj2/15-849/slides/12-tensor-rematerialization.pdf:image=https://www.cs.cmu.edu/~zhihaoj2/15-849/slides/12-tensor-rematerialization.pdf

github: https://github.com/dmlc/mxnet-memonger

背景

  • 深層学習モデルの学習にはパラメタ保存の何倍ものメモリが必要

どんなもの?

  • 訓練時(forward/backward)に必要なメモリ用を減らす方法の提案

先行研究と比べてどこがすごい?

  • 学習中における計算とメモリのトレードオフを考慮できる
  • (少なくともCNN, RNNを含む)一般的なモデルに適用できる

技術や手法のキモはどこ?

  • sec3: backward時に、ノード(演算)ごとに、メモリを確保するか既に不要なメモリを再利用(上書き)するかを シンプルなラベル付グラフアルゴリズムにより求める

(赤いところはメモリの共有ができる)

  • sec4: モデルを分割し、分割されたブロックの最初の層の活性状態量(隠れ層の出力)を保存する
    • 保存するデータ量が減少
    • backward時には再びforwardを行って活性化状態を復元する(追加の計算コストあり)

(シンプルでメモリ削減効果が高い)

どうやって有効だと検証した?

  • 全ての活性状態量、特徴マップを保存する場合
  • sec3の方法を用いてメモリ使用量削減を試みた場合
  • sec3+4の方法(checkpint)で計算コストと引き換えにさらにメモリ削減を試みた場合 で学習させてメモリ使用量を比較 層数1000のLSTMだとcheckpointで約256GB → 2GB以下に削減できた、計算時間は30%増

議論はある?

  • checkpoint作成箇所によって、効率が変わってくることはあるのか?
  • backward時にcheckpointの復元と勾配計算を同時にやることで、計算効率は少しは良くなるはず(スケジュールが難しいか?)