NNの学習で必要なGPUメモリを算出する

最近、ようやくGPUを用いたDeepLearningを仕事で使うようになりました。

しかし使ってみると頻繁に「out of memory exception」が発生したため、

今回はGPUのメモリ使用量について調べてみました。

GPUメモリ使用量算出方法

より具体的には、ニューラルネットワークの学習に必要なメモリ量を算出します。

結論から言うと、下記の式で算出できます。

必要メモリ量(byte) =

　(ニューロンの数 × Batch Size ＋ パラメータ数) × 2 (data & grad) × 4(byte)

これは、sonyのNeural Network Consoleのチュートリアルに掲載されていたものを参考にしました。

チュートリアル：最適な学習実行環境の選び方 – Docs - Neural Network Console

ドキュメントでは、Neural Network Consoleの使い方を提示しながら、主要な機能について説明します。セットアップから学習の開始までドキュメントを見て進めてみましょう。

support.dl.sony.com

CNNで用いられるメモリ量を計算する

私が実際に算出した例をご紹介します。

今回は、Kerasで使えるVGG16をそのまま使ってみます。

Kerasでモデルを作成し、model.summary()を実行すると、下記の出力が得られます。

デフォルトでは、画像サイズは224×224ですね。

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_3 (InputLayer)         (None, 224, 224, 3)       0         
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 224, 224, 64)      1792      
_________________________________________________________________
・
・
省略
・
・
_________________________________________________________________
flatten (Flatten)            (None, 25088)             0         
_________________________________________________________________
fc1 (Dense)                  (None, 4096)              102764544 
_________________________________________________________________
fc2 (Dense)                  (None, 4096)              16781312  
_________________________________________________________________
predictions (Dense)          (None, 1000)              4097000   
=================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
_________________________________________________________________

ニューロン数とパラメータ数

メモリ量算出のためにニューロン数とパラメータ数を知る必要があります。

パラメータ数に関しては、model.summary()で出力されるTotal paramを参照すればよいです。

Total params: 138,357,544 
Trainable params: 138,357,544 
Non-trainable params: 0

今回のケースだと、138,357,544がパラメータ数になります。

続いてニューロン数ですが、これはmodel.summary()のOutput Shapeから算出する必要があります。

算出の方法は、各レイヤーのニューロン数をすべて足す、です。

_________________________________________________________________
Output Shape              
=================================================================
(None, 224, 224, 3)   224 × 224 × 3 = 150,528    
_________________________________________________________________
(None, 224, 224, 64)  224 × 224 × 64 = 3,211,264 
_________________________________________________________________
・
・
省略
・
・
_________________________________________________________________
(None, 4096)                           4096
_________________________________________________________________
(None, 1000)                           1000
=================================================================
             150,528 + 3,211,264 + ... + 4096 + 1000 = 15,262,696

これを計算すると、15,262,696となります。

VGG16のような画像認識を目的とするようなCNNの場合、

ニューロン数は画像のサイズに依存します。

今回は224×224でしたが、もっと大きな画像を扱おうと思うとニューロン数も増えていきます。

Batch Size

Batch Sizeは自分で決める値です。今回はBatch Size＝8で計算してみます。

メモリ量算出

先ほどの式にニューロン数、パラメータ数、Batch Sizeを当てはめてみると、

必要メモリ量(byte) =

　(15,262,696 × 8 ＋ 138,357,544) × 2 (data & grad) × 4(byte)

= 2,083,672,896(byte) ≒ 2GB

となります。

VGG16を使って224×224サイズの画像をBatch Size=8で学習するには、

GPUメモリが約2GB必要になることがわかりました。

GPUメモリ内で処理するには

メモリ算出の式を見てもらえばわかりますが、特にCNNの場合、消費されるメモリは

ニューロン数(≒画像サイズ)とBatch Sizeが大きく効いてきます。

例えば、先ほどのVGG16をBatch Size=32にすると、メモリは4倍の約8GB必要になります。

また、画像サイズを512×512にすると、メモリは約6GB必要になります。

2018年11月に発売されたGeforce RTX 2080Tiでもメモリ容量は11GBなので

画像サイズとBatch Sizeに気を付けないとすぐにメモリ不足のエラーになってしまいます。

逆に最初からメモリの計算をしていれば無駄な時間を使わずに済みますね。

まとめ

ニューラルネットワークで使われるメモリ量の算出方法をまとめてみました。

ちなみに、今回記した方法が本当に正しいか、実は証明していません。（笑）

間違っていればご指摘お願いします。

また、実際にはGPUメモリをすべてニューラルネットワークに使えるわけではないので

今回の計算方法が正しくても少しマージンを取る必要があるのは間違いありません。

また、私はKerasでmodel.summary()で出力できるモデルしかメモリ算出できませんので

他に素晴らしい方法を知っている方はぜひ教えてください。

以上！！