ディープラーニング実行時のメモリ計算方法
ディープラーニングでGPU実行しているとき、メモリが足らずよくプロセスを落としたりします。
実際、事前にメモリをどのように計算するのだろうとふと思った時に、パッと出てこなかったので、ここにメモします。ほぼ参考の日本語訳です。
ディープラーニングモデルの層数
傾向としては深くなっている傾向にあり、年々GPUメモリの利用容量は拡大しています。画像認識コンペティションILSVRCでの結果を振り返るとよく分かります。
| ネットワーク | 年 | 層数 | Top-5 |
|---|---|---|---|
| Alexne | 2012 | 7 | 17.0 |
| VGG-19 | 2014 | 19 | 9.35 |
| GoogleNet | 2014 | 22 | 9.15 |
| Resnet-50 | 2015 | 50 | 6.71 |
| Resnet-152 | 2015 | 152 | 5.71 |
こうしてみると、Resnetに導入されたbatch normalization偉大ですね。。
なにがメモリを使うのか
ニューラルネットワークのパラメータ数は増えていく傾向で、今後も精度を求める限りは増え続けていくと思われます。パラメータとしては、例えば、ネットワークの訓練時では以下のパラメータがGPUメモリにのる必要があります。
１．中間層の出力
２．重みやバイアスといったパラメータ
３．バックプロパゲーション時の各素子での誤差
４．勾配法（SGD、Adamとか）のパラメータ
５．その他オーバーヘッドや最適化手法に依存するパラメータ（モーメンタムとか）
ネットワークで評価する際は学習時のみに用いるパラメータを除いた以下のパラメータとなります。
１．中間層の出力
２．重みやバイアスといったパラメータ
これらのパラメータのサイズがGPUメモリにのるようにネットワークは設計する必要があります。
GeForce GTX 1080 Tiで11GBのメモリを搭載しており、環境はどんどんよくなってきているものの、大きい画像サイズを扱う際やすべて畳み込みのネットワークを構成するときは注意が必要となります。
どうやって必要なメモリを計算するのか
例えば、畳込み層だと、重みの数は入力サイズに依存せず、入力チャネル数、カーネルサイズ、出力チャネル数を用いて重みの数を計算します。
具体的には、
出力チャネル数×（カーネルサイズ×カーネルサイズ）×入力チャネル数
となります。
以下のような畳込みニューラルネットワークに28x28の入力画像に適用したとします。
第3層：2x2のマックスプーリング
第4層：128素子の全結合層
第5層：10素子の全結合層
出力層：ソフトマックス
※畳込み層では入力サイズとフィーチャーマップサイズを同一となるようpadding=1、stride=1とします
この第1層での重みの数は
32x(3x3)×1=288個
となります。
また、各畳込みフィルタにはバイアスがつくので、バイアスの数は、
32x1=32個
となります。
計320個のパラメータとなり、これがカラー画像なので3チャネルの入力だと、960個になります。
カーネルサイズ3x3、出力チャネル32の畳込み層をもう一度かけた場合（第2層）、パラメータ数は以下となります。
第1層と合わせて、9,568個のパラメータが必要となります。
第3層のプーリング層ではマックスプーリングには重みもバイアスもないため、パラメータフリーな層となります。2x2のマックスプーリングを行い、サイズを28x28->14->14へとダウンサンプリングします。
第4層の全結合層では、以下のパラメータ数があります。
重み：出力数×入力数
バイアス：出力数
出力数は全結合層の素子数、入力数はマックスプーリング層の出力サイズとチャネル数となります。
よって重みの数は、
128x(14x14x32)=802,816
となり、これにバイアスの数128を加えて、全結合層の802944個がパラメータ数となります。
第5層でも同様で、第5層のパラメータ数は
10x128 + 10 = 1,290
となります。
まとめるとパラメータ数は、
第1層：320
第2層：9,248
第3層：なし
第4層：802,944
第5層：1,290
各層の総和：813,802個
となり、モデルサイズは以下となります。
813,802x(32/8)/(1024)^2 ~ 3.1MB
ここで、ニューラルネットワークのパラメータは32-bit floatsで表現されることとしました。
これらのパラメータに加えて、中間層の出力サイズを各層ごとに計算すると、
第1層：32x(28x28) = 25,088
第2層：32x(28x28) = 25,088
第3層：32x(14x14) = 6,272
第4層：128
第5層：10
各層の総和：56,586個
のパラメータがあります。
56,586x(32/8)/(1024)^2 ~ 0.2MB
つまり、評価する際は、3.3MBのメモリがGPUに必要となります。
訓練する際は、勾配の誤差等を考慮する必要があります。
概ねモデルサイズを3倍、中間層の出力サイズを2倍にします。
モデルサイズを3倍するのは、モーメンタムとかステップキャッシュとかを考慮してかけてます。
中間層を2倍にするのは、activationsの勾配を考慮してかける。勾配はactivationsの数と同一のため、2倍となります。
3.1MBx3 + 0.2MBx2 = 10MB
さらに、バッチサイズも考慮する必要があります。
上記の計算は1枚の画像入力に対しての計算であり、バッチサイズが128であったら、activationsの数にさらにバッチサイズをかける必要がある。バッチサイズが128であれば以下となります。
3.1MBx3 + 0.2MBx2x128 = 60.5MB
バッチサイズがクリティカルにメモリサイズへと影響を与えることがわかります。
VGG16の場合、どれぐらいメモリをつかうのか
によると、以下のとおりです。
INPUT: [224x224x3] memory: 224*224*3=150K weights: 0
CONV3-64: [224x224x64] memory: 224*224*64=3.2M weights: (3*3*3)*64 = 1,728
CONV3-64: [224x224x64] memory: 224*224*64=3.2M weights: (3*3*64)*64 = 36,864
POOL2: [112x112x64] memory: 112*112*64=800K weights: 0
CONV3-128: [112x112x128] memory: 112*112*128=1.6M weights: (3*3*64)*128 = 73,728
CONV3-128: [112x112x128] memory: 112*112*128=1.6M weights: (3*3*128)*128 = 147,456
POOL2: [56x56x128] memory: 56*56*128=400K weights: 0
CONV3-256: [56x56x256] memory: 56*56*256=800K weights: (3*3*128)*256 = 294,912
CONV3-256: [56x56x256] memory: 56*56*256=800K weights: (3*3*256)*256 = 589,824
CONV3-256: [56x56x256] memory: 56*56*256=800K weights: (3*3*256)*256 = 589,824
POOL2: [28x28x256] memory: 28*28*256=200K weights: 0
CONV3-512: [28x28x512] memory: 28*28*512=400K weights: (3*3*256)*512 = 1,179,648
CONV3-512: [28x28x512] memory: 28*28*512=400K weights: (3*3*512)*512 = 2,359,296
CONV3-512: [28x28x512] memory: 28*28*512=400K weights: (3*3*512)*512 = 2,359,296
POOL2: [14x14x512] memory: 14*14*512=100K weights: 0
CONV3-512: [14x14x512] memory: 14*14*512=100K weights: (3*3*512)*512 = 2,359,296
CONV3-512: [14x14x512] memory: 14*14*512=100K weights: (3*3*512)*512 = 2,359,296
CONV3-512: [14x14x512] memory: 14*14*512=100K weights: (3*3*512)*512 = 2,359,296
POOL2: [7x7x512] memory: 7*7*512=25K weights: 0
FC: [1x1x4096] memory: 4096 weights: 7*7*512*4096 = 102,760,448
FC: [1x1x4096] memory: 4096 weights: 4096*4096 = 16,777,216
FC: [1x1x1000] memory: 1000 weights: 4096*1000 = 4,096,000
TOTAL memory: 24M * 4 bytes ~= 93MB / image (only forward! ~*2 for bwd)
TOTAL params: 138M parameters
こういうのって、ある程度推論できるので、大きなデータを扱う際は事前チェックするとパラメータの調整がしやすくなります。
参考URL
CS231n Convolutional Neural Networks for Visual Recognition
http://imatge-upc.github.io/telecombcn-2016-dlcv/slides/D2L1-memory.pdf
