[論文レビュー] Memory-Efficient Implementation of DenseNets

Geoff Pleiss, Danlu Chen+4 more|arXiv (Cornell University)|Jul 21, 2017

Advanced Neural Network Applications参考文献 12被引用数 86

ひとこと要約

本論文は、特徴マップのメモリを二乗から線形へ削減することで、再計算と組み合わせたメモリ共有戦略を提案し、数百層にも及ぶDenseNetsを比較的少ない時間オーバーヘッドで訓練可能にする。

ABSTRACT

The DenseNet architecture is highly computationally efficient as a result of feature reuse. However, a naive DenseNet implementation can require a significant amount of GPU memory: If not properly managed, pre-activation batch normalization and contiguous convolution operations can produce feature maps that grow quadratically with network depth. In this technical report, we introduce strategies to reduce the memory consumption of DenseNets during training. By strategically using shared memory allocations, we reduce the memory cost for storing feature maps from quadratic to linear. Without the GPU memory bottleneck, it is now possible to train extremely deep DenseNets. Networks with 14M parameters can be trained on a single GPU, up from 4M. A 264-layer DenseNet (73M parameters), which previously would have been infeasible to train, can now be trained on a single workstation with 8 NVIDIA Tesla M40 GPUs. On the ImageNet ILSVRC classification dataset, this large DenseNet obtains a state-of-the-art single-crop top-1 error of 20.26%.

研究の動機と目的

• 高容量モデルのためのDenseNetのパラメータ効率と特徴の再利用を動機づける。
• 標準的なDenseNet訓練における二乗オーダーのメモリボトルネックを特定する。
• 訓練メモリを二乗から線形へ削減するメモリ共有戦略を提案する。
• メモリ予算下で非常に深いDenseNetsの訓練と競争力のあるImageNet性能を実証する。

提案手法

• DenseNetsにおける二乗オーダーのメモリの二つの源を特定する：前活性化バッチ正規化と連続結合。
• 連結出力用のShared Memory Storage 1とバッチ正規化出力用のShared Memory Storage 2を導入する。
• 逆伝搬中に連結とバッチ正規化を再計算して共有ストレージを充填し、すべての中間結果を保存する代わりにする。
• 勾配ストレージを層間で共有し、勾配の二乗成長を回避する。
• メモリと時間のオーバーヘッドを測定し、大幅なメモリ節約を達成しつつ、訓練時間は約15-20%の追加となることを示す。

実験結果

リサーチクエスチョン

• RQ1 共有ストレージの再利用と再計算によって、DenseNetsは低メモリ環境でも効果的に訓練できるか。
• RQ2 どれくらいのメモリを削減できるか（二乗から線形へ）と、それに伴う計算コストはいくらか。
• RQ3 メモリ効率の DenseNetsでImageNet上で達成可能な深さとパラメータ数の実用的限界は何か。

主な発見

• 提案された共有メモリ戦略により、メモリ消費は深さに対して線形となる。
• LuaTorchでは、160層モデルがナイーブ実装の約22%のメモリを使用し、12 GB予算内で約340層モデルの訓練を可能にする。
• PyTorchでは、非常に深いDenseNetsを単一GPUでほぼ500層まで訓練することが可能。
• 効率的な実装で訓練したDenseNetsは、264層（k=48、73Mパラメータ）でImageNetのtop-1エラー20.26%を達成。
• 最も深いCosine DenseNetはtop-1エラー20.26%、従来の最先端を上回る。
• 勾配ストレージを共有することは時間コストなしで有益であり；共有BN/連結ストレージの追加は約15-20%の時間オーバーヘッドを追加する。